JUPYTER: create dictionary in module file separate from notebook

use simple append+replace Jupyter API

Signed-off-by: Lance-Drane <[email protected]>

Author: Lance-Drane

examples-proposed/004-time-loop/README.md

@@ -21,3 +21,11 @@ To run the code, run:
 ```bash
 ./run.sh
 ```
+
+By default, this example will always _append_ a state file. If you prefer to see an example of how to _replace_ a state file, run:
+
+```bash
+EXAMPLE_REPLACE=1 ./run.sh
+```
+
+There is also a script `run-delayed.sh` which you can use instead of `run.sh` if you would like to simulate a delay between monitor steps.

examples-proposed/004-time-loop/mymodule/components.py

@@ -8,11 +8,15 @@
 from ipsframework import Component
 
 DELAY = bool(os.environ.get('EXAMPLE_DELAY'))
+REPLACE = bool(os.environ.get('EXAMPLE_REPLACE'))
 
+# templates are existing files from the input directory
+# names are what the notebook and the associated data file will be labeled with (you can leave off the .ipynb / .py)
 NOTEBOOK_1_TEMPLATE = 'basic.ipynb'
 NOTEBOOK_1_NAME = 'basic.ipynb'
 NOTEBOOK_2_TEMPLATE = 'bokeh-plots.ipynb'
 NOTEBOOK_2_NAME = 'bokeh-plots.ipynb'
+DATA_MODULE_NAME = 'data_files'
 
 
 class Init(Component):
@@ -37,10 +41,12 @@ def step(self, timestamp=0.0):
         self.services.initialize_jupyter_notebook(
             dest_notebook_name=NOTEBOOK_1_NAME,  # path is relative to JupyterHub directory
             source_notebook_path=NOTEBOOK_1_TEMPLATE,  # path is relative to input directory
+            data_module_name=DATA_MODULE_NAME,
         )
         self.services.initialize_jupyter_notebook(
             dest_notebook_name=NOTEBOOK_2_NAME,  # path is relative to JupyterHub directory
             source_notebook_path=NOTEBOOK_2_TEMPLATE,  # path is relative to input directory
+            data_module_name=DATA_MODULE_NAME,
         )
 
         # The time loop is configured in its own section of sim.conf
@@ -105,9 +111,16 @@ def step(self, timestamp=0.0, **keywords):
         with open(state_file, 'rb') as f:
             data = f.read()
 
-        # stage the state file in the JupyterHub directory
-        self.services.add_data_file_to_notebook(state_file, timestamp, NOTEBOOK_1_NAME)
-        self.services.add_data_file_to_notebook(state_file, timestamp, NOTEBOOK_2_NAME)
+        # stage the state file in the JupyterHub directory and update the module file to handle it
+        if REPLACE:
+            self.services.add_analysis_data_file(state_file, os.path.basename(state_file), DATA_MODULE_NAME, replace=True)
+        else:
+            self.services.add_analysis_data_file(
+                state_file,
+                f'{timestamp}_{os.path.basename(state_file)}',
+                DATA_MODULE_NAME,
+                timestamp=timestamp,
+            )
 
         print('SEND PORTAL DATA', timestamp, data, file=stderr)
         self.services.send_portal_data(timestamp, data)

examples-proposed/004-time-loop/sim/input_dir/basic.ipynb

@@ -13,7 +13,7 @@
     "# In this example, this notebook is generated during the time loop.\n",
     "\n",
     "mapping = {}\n",
-    "for file in DATA_FILES:\n",
+    "for _timestep, file in DATA_FILES:\n",
     "    with open(file, 'rb') as f:\n",
     "        mapping[file] = f.read()\n",
     "print(mapping)\n"

examples-proposed/004-time-loop/sim/input_dir/bokeh-plots.ipynb

@@ -21,10 +21,10 @@
     "\n",
     "DATA = []\n",
     "# create DATA list, will depend on user input type (i.e. 'hdf5', 'json')\n",
-    "for file in DATA_FILES:\n",
+    "for _timestep, file in DATA_FILES:\n",
     "    with open(file, 'rb') as f:\n",
     "        DATA.append(json.load(f))\n",
-    "x = [float(f.rpartition('/')[2]) for f in DATA_FILES]\n",
+    "x = list(DATA_FILES.keys())\n",
     "\n",
     "COLORS = ['red', 'green', 'blue']\n",
     "\n",
@@ -51,7 +51,9 @@
     "    for idx, prop in enumerate(paths):\n",
     "        y = [get_data(d, prop) for d in DATA]\n",
     "        graph.line(x, y, line_color=COLORS[idx % len(COLORS)], line_dash='solid', legend_label='_'.join(prop))\n",
-    "    show(graph)\n"
+    "    show(graph)\n",
+    "\n",
+    "# TODO - add cell at end which shows history as a simulation (do it frame by frame instead of a line graph)\n"
    ]
   }
  ],

ipsframework/jupyter.py

@@ -1,4 +1,8 @@
 """
+NOTE: this is not intended to be a public API for framework users, use instead:
+  - "services.initialize_jupyter_notebook" (to set up the JupyterHub interaction for a notebook, done only once)
+  - "services.add_analysis_data_file" (each time you want to add or remove a data file from JupyterHub)
+
 This module is designed to help generate JupyterNotebooks to be used with IPS Portal analysis.
 Some parts of the script will need direction from users on the Framework side to generate.
 
@@ -10,18 +14,13 @@
 ...in a shell on Jupyter NERSC.
 """
 
+import re
 from os.path import sep
 from pathlib import Path
 from typing import Optional
 
 import nbformat as nbf
 
-HOOK = '### This cell autogenerated by IPS Framework. DO NOT EDIT UNTIL IPS RUN IS FINALIZED. ###'
-"""This hook is used to determine which "cell" the IPS framework should work with.
-
-It is written to a notebook cell on initializing it, and is searched for when adding a data file to it.
-"""
-
 DIRECTORY_VARIABLE_NAME = 'DATA_DIR'
 
 
@@ -34,149 +33,104 @@ def replace_last(source_string: str, old: str, new: str) -> str:
     return f'{head}{new}{tail}'
 
 
-def _initial_jupyter_file_notebook_cell(dest: str, files_variable_name: str) -> str:
-    return f"""{HOOK}
+def _initial_data_file_code(dest: str, files_variable_name: str) -> str:
+    return f"""# This file should be imported by a jupyter notebook. DO NOT EDIT UNTIL IPS RUN IS FINALIZED> 
 
 import os
 
 # NOTE: directory should be sim_name plus the run id from the Portal
 {DIRECTORY_VARIABLE_NAME} = '{str(Path(dest).parent / 'data') + sep}'
-# Uncomment below line to implicitly use any state files saved in the data directory, note that the IPS framework explicitly lists out each file used
-#{files_variable_name} = os.listdir('data')
-# files created during the run
-{files_variable_name} = [
-]
+{files_variable_name} = {{
+}}
 """
 
 
-def initialize_jupyter_notebook(dest: str, src: str, variable_name: str, index: int):
+def initialize_jupyter_notebook(notebook_dest: str, notebook_src: str, module_name: str, variable_name: str, index: int):
     """Create a new notebook from an old notebook, copying the result from 'src' to 'dest'.
 
+    This adds an additional cell which will import the data files. The notebook should not be written again after this function.
+
     Params:
-      - dest - location of notebook to create on filesystem (absolute file path)
-      - src - location of source notebook on filesystem (is not overwritten unless src == dest)
+      - notebook_dest - location of notebook to create on filesystem (absolute file path)
+      - notebook_src - location of source notebook on filesystem (is not overwritten unless src == dest)
+      - module_name - name of the python module which will contain the data file list
       - variable_name: what to call the variable
       - index: insert new cells at position before this value (will not remove preexisting cells)
-      - initial_data_files: optional list of files to initialize the notebook with
-
     """
     # to avoid conversion, use as_version=nbf.NO_CONVERT
-    nb: nbf.NotebookNode = nbf.read(src, as_version=4)
+    nb: nbf.NotebookNode = nbf.read(notebook_src, as_version=4)
 
     nb['cells'] = (
-        # warning notification for users inspecting the file, unused programatically
-        [nbf.v4.new_markdown_cell('# WARNING: Do not manually modify this file until the IPS simulation is complete.')]
-        + nb['cells'][:index]
+        nb['cells'][:index]
         + [
             # explicitly mark the IPS cell for users inspecting the file, unused programatically
             nbf.v4.new_markdown_cell('## Next cell generated by IPS Framework'),
-            nbf.v4.new_code_cell(_initial_jupyter_file_notebook_cell(dest, variable_name)),
+            nbf.v4.new_code_cell(f"""
+from {module_name} import {variable_name}
+import importlib
+
+importlib.reload('{variable_name}')
+"""),
         ]
         + nb['cells'][index:]
     )
 
     nbf.validate(nb)
-    with open(dest, 'w') as f:
+    with open(notebook_dest, 'w') as f:
         nbf.write(nb, f)
 
 
-def add_data_file_to_notebook(dest: str, data_file: str, index: Optional[int] = None):
-    """Add data file to notebook list.
+def initialize_jupyter_import_module_file(dest: str, variable_name: str):
+    """Create a new notebook from an old notebook, copying the result from 'src' to 'dest'.
 
     Params:
-      - dest: path to notebook which will be modified
-      - data_file: data file we add to the notebook
-      - index: optional index of the IPS notebook cell. If not provided, search through the notebook via an expected string hook.
+      - dest - location of notebook to create on filesystem (absolute file path)
+      - variable_name: what to call the variable
     """
-    nb: nbf.NotebookNode = nbf.read(dest, as_version=4)
-    if index is None:
-        index = next((i for i, e in enumerate(nb['cells']) if HOOK in e['source']), -1)
-    if index < 0:
-        raise Exception('Cannot find IPS notebook node')
-    ips_cell: str = nb['cells'][index]['source']
-
-    if ips_cell.find(f"f'{{{DIRECTORY_VARIABLE_NAME}}}{data_file}',\n]") != -1:
-        # The data file is already referenced in the notebook, so there's nothing else to do
-        return
-
-    # data file does not exist, so we need to add it
-    # search from right of string for the ']' character, should work assuming user does not modify the cell past the variable definition
-    result = replace_last(ips_cell, ']', f"f'{{{DIRECTORY_VARIABLE_NAME}}}{data_file}',\n]")
-    nb['cells'][index]['source'] = result
 
     with open(dest, 'w') as f:
-        nbf.write(nb, f)
-
+        f.write(_initial_data_file_code(dest, variable_name))
 
-def remove_data_file_from_notebook(dest: str, data_file: str, index: Optional[int] = None):
-    """Remove a specific data file from the notebook list.
 
-    Params:
-      - dest: path to notebook which will be modified
-      - data_file: data file we remove from the notebook
-      - index: optional index of the IPS notebook cell. If not provided, search through the notebook via an expected string hook.
+def update_module_file_with_data_file(dest: str, data_file: str, replace: bool, timestamp: float = 0.0) -> Optional[str]:
     """
-    nb: nbf.NotebookNode = nbf.read(dest, as_version=4)
-    if index is None:
-        index = next((i for i, e in enumerate(nb['cells']) if HOOK in e['source']), -1)
-    if index < 0:
-        raise Exception('Cannot find IPS notebook node')
-    ips_cell: str = nb['cells'][index]['source']
-
-    head, sep, tail = ips_cell.rpartition(f"f'{{{DIRECTORY_VARIABLE_NAME}}}{data_file}',\n")
-    if sep == '':
-        # existing match not found, so there's nothing left to remove
-        return
-    result = f'{head}\n{tail}'
-    nb['cells'][index]['source'] = result
-
-    with open(dest, 'w') as f:
-        nbf.write(nb, f)
-
-
-def remove_last_data_file_from_notebook(dest: str, index: Optional[int] = None) -> Optional[str]:
-    """Obtain the last data file entry in a notebook, remove it, and then return the name of the file.
-
-    Note that this function assumes the notebook maintains a specific format.
+    Params:
+      - dest: path to module file which will be modified
+      - data_file: file which will be added to the module
+      - replace: if True, we can update
+      - timestamp: key we associate the data file with
 
     Returns:
-      - None if there were no data entries in the notebook, the name of the file removed (without the directory) as a string if there was
+      - if we replaced a file, the name of the file which was replaced; otherwise, None
     """
-    nb: nbf.NotebookNode = nbf.read(dest, as_version=4)
-    if index is None:
-        index = next((i for i, e in enumerate(nb['cells']) if HOOK in e['source']), -1)
-    if index < 0:
-        raise Exception('Cannot find IPS notebook node')
-    ips_cell: str = nb['cells'][index]['source']
-
-    search_hook = f"f'{{{DIRECTORY_VARIABLE_NAME}}}"
-
-    start_index = ips_cell.rfind(search_hook)
-    if start_index == -1:
-        # no data files have been added, nothing to do
-        return None
-
-    ret = None
-    file_name_start_index = start_index + len(search_hook)
-    end_index = file_name_start_index
-    while True:
-        try:
-            end_char = ips_cell[end_index]
-            end_index += 1
-            if end_char == '\n':
-                # each entry gets its own "line", so we don't need to search anymore
-                break
-            if ips_cell[end_index] == "'" and ips_cell[end_index - 1] != '\\':
-                # we have found the name of the file
-                ret = ips_cell[file_name_start_index:end_index]
-        except IndexError:
-            # improperly formatted file (reached EOF), fall back to just removing everything after the break
-            return None
-
-    result = ips_cell[:start_index] + ips_cell[end_index:]
-    nb['cells'][index]['source'] = result
+    with open(dest, 'r') as f:
+        old_module_code = f.read()
+
+    replaced_file_name = None
+
+    timestamp_regex = str(timestamp).replace('.', '\\.')
+    directory_str = '\{' + DIRECTORY_VARIABLE_NAME + '\}'
+
+    search_pattern = f"{timestamp_regex}: f'{directory_str}(.*)',"
+
+    found_match = re.search(search_pattern, old_module_code)
+    if found_match:  # timestamp already exists
+        if replace:
+            replaced_file_name = found_match.group(1)
+            if replaced_file_name == data_file:
+                # in this case, we're not actually removing an obsolete file, so no need to write to the module file
+                # return None because we've already directly replaced the file
+                return None
+            new_module_code = re.sub(search_pattern, f"{timestamp}: f'{{{DIRECTORY_VARIABLE_NAME}}}{data_file}',", old_module_code)
+        else:
+            raise ValueError(
+                f"For timestamp entry {timestamp}, you are trying to replace '{found_match.group(1)}' with '{data_file}' . If this was intended, you must explicitly set 'replace=True' on the IPS function call."
+            )
+    else:  # timestamp does not exist, so add it
+        # search from right of string for the '}' character, should work assuming user does not modify the cell past the variable definition
+        new_module_code = replace_last(old_module_code, '}', f"{timestamp}: f'{{{DIRECTORY_VARIABLE_NAME}}}{data_file}',\n" + '}')
 
     with open(dest, 'w') as f:
-        nbf.write(nb, f)
+        f.write(new_module_code)
 
-    return ret
+    return replaced_file_name