Implementing new API for NDK (#122)

* Implementing new API for NDK * Updating gallery examples to use the new API * Updating 3d scenario result * Cleaning up scenario __init__ * Updating api docs * Compiling problem explicitly * Small fixes to plot full scenario * Creating NDK's own Problem class * Removing 'scenario_id' attribute * Cleaning up * Fixing docs * Decoupling render layout from compile problem * Moving grid creation to its own class * Moving center_frequency to the scenario root * Cleaning up * Cleaning up code. Invalidating grid and compiled problem when center_frequency is changed * Adapting documentation * Updating plot full scenario to use the procedural approach * Updating API docs adding the Grid and Problem pages * Fixing bug with 3D example * Improving docs, updating hash of 3d results file * Improving docs * Small improvements to notebook * Addressing review suggestion: replacing enum with subdirectory containing the built in scenarios * Update src/neurotechdevkit/grid.py Co-authored-by: charles <[email protected]> * Addressing review suggestion * Fixing issue with commited PR suggestion * Fixing bug * Removing layer_ids from problem creation. Removing need of creating material_layers, using material_masks keys instead * Fixing print * Replacing numpy arrays with plain lists for simple attributes like position, origin and direction * Refactoring * Refactoring * Fixing hash of results file * Simplifying API * Setting center_frequency in built_in scenarios * Adding SliceAxis to spellcheck's whitelist * Adding unit tests and cleaning up code * Adding predefined target options * Fix typo * Showing how the predefined targt options can be used * Improving plot full scenario masks creation * Addressing review suggestions * Fixing spellchecking error by removing unnecessary types from docstrings * Updating 3d file hash * update docstrings --------- Co-authored-by: charles <[email protected]> Co-authored-by: d-lucena <[email protected]>
agencyenterprise · Aug 9, 2023 · 7696119 · 7696119
1 parent 6e53d78
commit 7696119
Show file tree

Hide file tree

Showing 58 changed files with 1,565 additions and 1,711 deletions.
diff --git a/docs/api/grid.md b/docs/api/grid.md
@@ -0,0 +1,5 @@
+# Grid
+
+::: neurotechdevkit.grid
+    options:
+      show_root_toc_entry: false
diff --git a/docs/api/make.md b/docs/api/make.md
diff --git a/docs/api/materials.md b/docs/api/materials.md
@@ -0,0 +1,5 @@
+# Materials
+
+::: neurotechdevkit.materials
+    options:
+      show_root_toc_entry: false
diff --git a/docs/api/problem.md b/docs/api/problem.md
@@ -0,0 +1,5 @@
+# Problem
+
+::: neurotechdevkit.problem
+    options:
+      show_root_toc_entry: false
diff --git a/docs/api/scenarios.md b/docs/api/scenarios.md
@@ -12,35 +12,25 @@
     options:
       show_root_heading: true
 
-::: neurotechdevkit.scenarios.Scenario0
+::: neurotechdevkit.scenarios.built_in.Scenario0
     options:
       show_root_heading: true
-      members:
-        - _SCENARIO_ID
 
-::: neurotechdevkit.scenarios.Scenario1_2D
+::: neurotechdevkit.scenarios.built_in.Scenario1_2D
     options:
       show_root_heading: true
-      members:
-        - _SCENARIO_ID
 
-::: neurotechdevkit.scenarios.Scenario1_3D
+::: neurotechdevkit.scenarios.built_in.Scenario1_3D
     options:
       show_root_heading: true
-      members:
-        - _SCENARIO_ID
 
-::: neurotechdevkit.scenarios.Scenario2_2D
+::: neurotechdevkit.scenarios.built_in.Scenario2_2D
     options:
       show_root_heading: true
-      members:
-        - _SCENARIO_ID
 
-::: neurotechdevkit.scenarios.Scenario2_3D
+::: neurotechdevkit.scenarios.built_in.Scenario2_3D
     options:
       show_root_heading: true
-      members:
-        - _SCENARIO_ID
 
 ::: neurotechdevkit.scenarios.Target
     options:

diff --git a/docs/api/utils.md b/docs/api/utils.md
diff --git a/docs/examples/plot_3d.py b/docs/examples/plot_3d.py
@@ -19,8 +19,8 @@
 
 import neurotechdevkit as ndk
 
-URL = "https://neurotechdevkit.s3.us-west-2.amazonaws.com/result-scenario-2-3d.tz"
-known_hash = "6a5de26466028c673d253ca014c75c719467ec6c28d7178baf9287b44ad15191"
+URL = "https://neurotechdevkit.s3.us-west-2.amazonaws.com/result-scenario-2-3d-v3.tz"
+known_hash = "f9c1305fe4be4d348587fdcdfa54bc1a994e6e7dea866bc88ecc7875608c6010"
 downloaded_file_path = pooch.retrieve(url=URL, known_hash=known_hash, progressbar=True)
 result = ndk.load_result_from_disk(downloaded_file_path)
 

diff --git a/docs/examples/plot_adding_custom_source.py b/docs/examples/plot_adding_custom_source.py
@@ -27,20 +27,20 @@
 # - delay `(float, optional)`: the delay (in seconds) that the source will wait before
 #   emitting. Defaults to 0.0.
 
-import numpy as np
-
 import neurotechdevkit as ndk
 
 source = ndk.sources.FocusedSource2D(
-    position=np.array([0.00, 0.0]),
-    direction=np.array([0.9, 0.0]),
+    position=[0.00, 0.0],
+    direction=[0.9, 0.0],
     aperture=0.01,
     focal_length=0.01,
     num_points=1000,
 )
 
-scenario = ndk.make("scenario-0-v0")
-scenario.add_source(source)
+scenario = ndk.built_in.Scenario0()
+scenario.sources = [source]
+scenario.make_grid()
+scenario.compile_problem()
 result = scenario.simulate_steady_state()
 assert isinstance(result, ndk.results.SteadyStateResult2D)
 result.render_steady_state_amplitudes()

diff --git a/docs/examples/plot_customized_center_frequency.py b/docs/examples/plot_customized_center_frequency.py
@@ -11,19 +11,18 @@
 
 CENTER_FREQUENCY = 6e5
 
-scenario = ndk.make("scenario-0-v0")
-
-# using default material layers
-scenario.material_layers = ["water", "cortical_bone", "brain", "tumor"]
+scenario = ndk.scenarios.built_in.Scenario0()
 
 # Customizing material properties
 scenario.material_properties = {
     "tumor": ndk.materials.Material(
         vp=1850.0, rho=1250.0, alpha=0.8, render_color="#94332F"
     ),
 }
-
-result = scenario.simulate_steady_state(center_frequency=CENTER_FREQUENCY)
+scenario.center_frequency = CENTER_FREQUENCY
+scenario.make_grid()
+scenario.compile_problem()
+result = scenario.simulate_steady_state()
 assert isinstance(result, ndk.results.SteadyStateResult2D)
 result.render_steady_state_amplitudes(show_material_outlines=False)
 

diff --git a/docs/examples/plot_full_scenario.py b/docs/examples/plot_full_scenario.py
@@ -10,114 +10,62 @@
 The following code is a simplified implementation of NDK's Scenario 1.
 """
 # %%
-# ## Implementing a Scenario
+# Implementing a Scenario
 import numpy as np
-import stride
 
 from neurotechdevkit import sources
+from neurotechdevkit.grid import Grid
+from neurotechdevkit.problem import Problem
 from neurotechdevkit.results import SteadyStateResult2D
-from neurotechdevkit.scenarios import (
-    Scenario2D,
-    Target,
-    add_material_fields_to_problem,
-    make_grid,
+from neurotechdevkit.scenarios import Scenario2D, Target
+
+# %%
+# Creating the scenario
+scenario = Scenario2D()
+scenario.center_frequency = 5e5  # Hz
+scenario.target = Target(
+    target_id="target_1", center=[0.064, 0.0], radius=0.004, description=""
+)
+scenario.material_properties = {}
+scenario.origin = [0.0, -0.035]
+scenario.sources = [
+    sources.FocusedSource2D(
+        position=[0.0, 0.0],
+        direction=[1.0, 0.0],
+        aperture=0.064,
+        focal_length=0.064,
+        num_points=1000,
+    )
+]
+scenario.material_outline_upsample_factor = 8
+
+# %%
+# Creating grid
+grid = Grid.make_grid(
+    extent=(0.12, 0.07),  # m
+    speed_water=1500,
+    center_frequency=scenario.center_frequency,
+    ppw=6,
 )
 
+scenario.grid = grid
+
+# %%
+# Creating masks
+
+
+def fill_mask(mask, start, end, dx):
+    # fill linearly along the x axis
+    if end is None:
+        n = int(start / dx)
+        mask[n:] = True
+    else:
+        n = int(start / dx)
+        m = int(end / dx)
+        mask[n:m] = True
 
-class FullScenario(Scenario2D):
-    """This Scenario is based on benchmark 4 of the following paper:
-
-    Jean-Francois Aubry, Oscar Bates, Christian Boehm, et al., "Benchmark problems for
-    transcranial ultrasound simulation: Intercomparison of compressional wave models",
-    The Journal of the Acoustical Society of America 152, 1003 (2022);
-    doi: 10.1121/10.0013426
-    https://asa.scitation.org/doi/pdf/10.1121/10.0013426
-    """
-
-    _SCENARIO_ID = "the_id_for_this_scenario"
-
-    """
-    Target attributes:
-        target_id: the string id of the target.
-        center: the location of the center of the target (in meters).
-        radius: the radius of the target (in meters).
-        description: a text describing the target.
-    """
-    _TARGET_OPTIONS = {
-        "target_1": Target("target_1", np.array([0.064, 0.0]), 0.004, ""),
-    }
-
-    """
-    The order of returned materials defines the layering of the scenario.
-    """
-    material_layers = [
-        "water",
-        "skin",
-        "cortical_bone",
-        "trabecular_bone",
-        "brain",
-    ]
-
-    def __init__(self, complexity="fast"):
-        """
-        Instantiate a new scenario.
-        The origin defines the spatial coordinates of grid position (0, 0, 0).
-        """
-        self._target_id = "target_1"
-
-        super().__init__(
-            origin=np.array([0.0, -0.035]),
-            complexity=complexity,
-        )
-
-    @property
-    def _material_outline_upsample_factor(self) -> int:
-        """
-        The factor to use when upsampling the material field before
-        detecting transitions between materials. If the factor is N, then each pixel
-        will be split into N^2 pixels. Defaults to 1 (no resampling).
-        """
-        return 8
-
-    def _compile_problem(self, center_frequency) -> stride.Problem:
-        """The problem definition for the scenario."""
-        extent = np.array([0.12, 0.07])  # m
-        # scenario constants
-        speed_water = 1500  # m/s
-
-        # desired resolution for complexity=fast
-        ppw = 6
-
-        # compute resolution
-        dx = speed_water / center_frequency / ppw  # m
-
-        grid = make_grid(extent=extent, dx=dx)
-        problem = stride.Problem(
-            name=f"{self.scenario_id}-{self.complexity}", grid=grid
-        )
-        problem = add_material_fields_to_problem(
-            problem=problem,
-            materials=self.get_materials(center_frequency),
-            layer_ids=self.layer_ids,
-            masks={
-                name: _create_scenario_1_mask(name, problem.grid)
-                for name in self.material_layers
-            },
-        )
-        return problem
-
-    def get_default_source(self) -> sources.Source:
-        """The transducer source for the scenario."""
-        return sources.FocusedSource2D(
-            position=np.array([0.0, 0.0]),
-            direction=np.array([1.0, 0.0]),
-            aperture=0.064,
-            focal_length=0.064,
-            num_points=1000,
-        )
-
-
-def _create_scenario_1_mask(material, grid):
+
+def create_masks(grid):
     # layers are defined by X position
     dx = grid.space.spacing[0]
 
@@ -131,56 +79,44 @@ def _create_scenario_1_mask(material, grid):
             0.0835,  # brain
         ]
     )
+    layers = np.array(
+        ["water", "skin", "cortical_bone", "trabecular_bone", "cortical_bone", "brain"]
+    )
     interfaces = np.cumsum(layers_m)
 
-    mask = np.zeros(grid.space.shape, dtype=bool)
+    mask_materials = {}
 
-    if material == "water":
-        _fill_mask(mask, start=0, end=interfaces[0], dx=dx)
+    for material in np.unique(layers):
+        mask = np.zeros(grid.space.shape, dtype=bool)
+        for index in np.where(layers == material)[0]:
+            start = interfaces[index - 1] if material != "water" else 0
+            end = interfaces[index] if material != "brain" else None
+            fill_mask(mask, start=start, end=end, dx=dx)
+        mask_materials[material] = mask
+    return mask_materials
 
-    elif material == "skin":
-        _fill_mask(mask, start=interfaces[0], end=interfaces[1], dx=dx)
 
-    elif material == "cortical_bone":
-        _fill_mask(mask, start=interfaces[1], end=interfaces[2], dx=dx)
-        _fill_mask(mask, start=interfaces[3], end=interfaces[4], dx=dx)
-
-    elif material == "trabecular_bone":
-        _fill_mask(mask, start=interfaces[2], end=interfaces[3], dx=dx)
-
-    elif material == "brain":
-        _fill_mask(mask, start=interfaces[4], end=None, dx=dx)
-
-    else:
-        raise ValueError(material)
-
-    return mask
-
-
-def _fill_mask(mask, start, end, dx):
-    # fill linearly along the x axis
-    if end is None:
-        n = int(start / dx)
-        mask[n:] = True
-    else:
-        n = int(start / dx)
-        m = int(end / dx)
-        mask[n:m] = True
+scenario.material_masks = create_masks(grid)
 
 
 # %%
-# ## Creating the scenario
-scenario = FullScenario()
+# Rendering the layout
+scenario.render_layout()
 
 # %%
-# ## Rendering the scenario layout
-scenario.render_layout()
+# Creating problem
+
+problem = Problem(grid=grid)
+problem.add_material_fields(
+    materials=scenario.materials,
+    masks=scenario.material_masks,
+)
+scenario.problem = problem
 
 # %%
-# ## Running the scenario
+# Rendering the simulation
 result = scenario.simulate_steady_state()
 assert isinstance(result, SteadyStateResult2D)
 result.render_steady_state_amplitudes(show_material_outlines=False)
 
-
 # %%