Miscellaneous Changes

src/NanoParticleTools/builder.py

@@ -53,12 +53,14 @@
                          **kwargs)
         self.connect()
 
-    def get_grouped_docs(self) -> List[Dict]:
+    def get_grouped_docs(self, additional_keys: List = None) -> List[Dict]:
         group_keys = [
             "data.n_dopants", "data.n_dopant_sites", "data.formula",
             "data.nanostructure_size", "data.formula_by_constraint",
             "data.excitation_power", "data.excitation_wavelength"
         ]
+        if additional_keys is not None and isinstance(additional_keys, list):
+            group_keys += additional_keys
         return self.source.groupby(keys=group_keys,
                                    criteria=self.docs_filter,
                                    properties=["_id"])

src/NanoParticleTools/machine_learning/data/utils.py

@@ -5,11 +5,12 @@
 import os
 
 SUNSET_SPECIES_TABLE = {
-    1: ["Yb", "Er", "Mg"],
-    2: ["Yb", "Er"],
-    3: ["Yb", "Er", "Mg", "Tm"],
-    4: ["Yb", "Er"],
-    5: ["Yb", "Er", "Nd"]
+    1: sorted(["Yb", "Er", "Xsurfacesix"]),
+    2: sorted(["Yb", "Er"]),
+    3: sorted(["Yb", "Er", "Xsurfacesix", "Tm"]),
+    4: sorted(["Yb", "Er"]),
+    5: sorted(["Yb", "Er", "Nd"]),
+    6: sorted(['Yb', 'Er', "Xsurfacesix", 'Tm', 'Nd', 'Ho', 'Eu', 'Sm', 'Dy'])
 }
 
 

src/NanoParticleTools/machine_learning/modules/ensemble.py

@@ -28,7 +28,7 @@ def ensemble_forward(self, data: Data,
             output.append(y_hat)
 
         x = torch.cat(output, dim=-1)
-        return {'y': x, 'y_hat': x.mean(-1), 'std': x.std()}
+        return {'y': x, 'y_hat': x.mean(-1), 'std': x.std(-1)}
 
     def evaluate_step(self, data: Data) -> tuple[torch.Tensor, torch.Tensor]:
         output = []
@@ -52,6 +52,6 @@ def predict_step(
 
         x = torch.cat(output, dim=-1)
         if return_stats:
-            return {'y': x, 'y_hat': x.mean(-1), 'std': x.std()}
+            return {'y': x, 'y_hat': x.mean(-1), 'std': x.std(-1)}
         else:
             return x.mean(-1)

src/NanoParticleTools/optimization/callbacks.py

@@ -1,4 +1,4 @@
-from NanoParticleTools.util.visualization import plot_nanoparticle_from_arrays
+from NanoParticleTools.util.visualization import plot_nanoparticle
 from NanoParticleTools.machine_learning.data import FeatureProcessor
 
 from maggma.stores import Store
@@ -11,7 +11,8 @@
 from uuid import uuid4
 
 
-def get_plotting_fn(feature_processor: FeatureProcessor) -> Callable:
+def get_plotting_fn(feature_processor: FeatureProcessor,
+                    as_np_array: bool = False) -> Callable:
     n_elements = len(feature_processor.possible_elements)
 
     def plotting_fn(x, f=None, accept=None):
@@ -20,19 +21,30 @@ def plotting_fn(x, f=None, accept=None):
 
         plt.figure()
         n_constraints = len(x) // (n_elements + 1)
-        plot_nanoparticle_from_arrays(
+        fig = plot_nanoparticle(
             np.concatenate(([0], x[-n_constraints:])),
             x[:-n_constraints].reshape(n_constraints, -1),
-            dpi=80,
-            elements=feature_processor.possible_elements,
-        )
+            dpi=300,
+            elements=feature_processor.possible_elements)
         if f is not None:
             plt.text(0.1,
                      0.95,
                      f'UV Intensity={np.power(10, -f)-100:.2f}',
                      fontsize=20,
                      transform=plt.gca().transAxes)
-        return plt
+        if as_np_array:
+            # If we haven't already shown or saved the plot, then we need to
+            # draw the figure first.
+            fig.canvas.draw()
+            # Now we can save it to a numpy array.
+            data = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
+            data = data.reshape(fig.canvas.get_width_height()[::-1] + (3, ))
+
+            # Close the figure to remove it from the buffer
+            plt.close(fig)
+            return data
+        else:
+            return fig
 
     return plotting_fn
 

src/NanoParticleTools/optimization/scipy_optimize.py

@@ -12,7 +12,7 @@
 from collections.abc import Callable
 
 
-def get_bounds(n_constraints: int, n_elements: int) -> Bounds:
+def get_bounds(n_constraints: int, n_elements: int, **kwargs) -> Bounds:
     r"""
     Get the Bounds which are utilized by scipy minimize.
 
@@ -34,7 +34,8 @@ def get_bounds(n_constraints: int, n_elements: int) -> Bounds:
         (np.zeros(num_dopant_nodes), np.zeros(n_constraints)))
     max_bounds = np.concatenate(
         (np.ones(num_dopant_nodes), np.ones(n_constraints)))
-    bounds = Bounds(min_bounds, max_bounds)
+    min_bounds[-1] = 1
+    bounds = Bounds(min_bounds, max_bounds, **kwargs)
     return bounds
 
 

src/NanoParticleTools/species_data/species.py

@@ -84,6 +84,14 @@
 
     # The naming convention should always start with an X, since the symbol
     # cannot start with an existing element's symbol
+    LEGACY_SURFACE_NAMES = {
+        'Na': 'Surface',
+        'Al': 'Surface3',
+        'Si': 'Surface4',
+        'P': 'Surface5',
+        'Mg': 'Surface6',
+    }
+
     SURFACE_DOPANT_SYMBOLS_TO_NAMES = {
         'Xsurfaceone': 'Surface',
         'Xsurfacethree': 'Surface3',
@@ -117,7 +125,16 @@
     def __init__(self,
                  symbol: str,
                  molar_concentration: float,
-                 n_levels: int | None = None):
+                 n_levels: int | None = None,
+                 legacy_calc: bool = False):
+        self.legacy_calc = legacy_calc
+
+        if self.legacy_calc:
+            # If this is an older (legacy) calc, we need to convert the hacked
+            # surface species to the current.
+            if symbol in self.LEGACY_SURFACE_NAMES:
+                symbol = self.LEGACY_SURFACE_NAMES[symbol]
+
         if symbol in self.SURFACE_DOPANT_NAMES_TO_SYMBOLS:
             symbol = self.SURFACE_DOPANT_NAMES_TO_SYMBOLS[symbol]
         self.symbol = symbol

src/NanoParticleTools/util/visualization.py

@@ -16,125 +16,117 @@
 }
 
 
-def plot_nanoparticle_from_arrays(radii: np.array,
-                                  concentrations: np.array,
-                                  dpi=150,
-                                  as_np_array=False,
-                                  elements=['Yb', 'Er', 'Nd']):
+def plot_nanoparticle(radii: np.ndarray | list[NanoParticleConstraint],
+                      concentrations: np.array = None,
+                      dopant_specifications: list[tuple] = None,
+                      dpi=150,
+                      as_np_array=False,
+                      elements=['Yb', 'Er', 'Nd'],
+                      ax: plt.Axes = None,
+                      emissions: float = None):
     if 'Y' not in elements:
+        # Add Y, the host element
         elements = elements + ['Y']
 
-    # Fill in the concentrations with Y
-    concentrations_with_y = np.concatenate(
-        (concentrations, 1 - concentrations.sum(axis=1, keepdims=True)),
-        axis=1)
-
+    if isinstance(radii[0], NanoParticleConstraint):
+        # Convert this to an array
+        radii = np.array([0] + [c.radius for c in radii])
+    if not isinstance(radii, np.ndarray):
+        # If it is a list, it is already in the format we require
+        raise TypeError(
+            'radii should be an array of radii or list of contraints')
+
+    if concentrations is None and dopant_specifications is None:
+        raise RuntimeError(
+            'Must specify one of concentrations or dopant specifications')
+    elif dopant_specifications is not None:
+        # convert this to an array
+        n_layers = len(radii) - 1
+        dopant_dict = [{key: 0 for key in elements} for _ in range(n_layers)]
+        for dopant in dopant_specifications:
+            dopant_dict[dopant[0]][dopant[2]] = dopant[1]
+
+        # Fill in the rest with 'Y'
+        for layer in dopant_dict:
+            layer['Y'] = 1 - sum(layer.values())
+
+        vals = [[layer[el] for el in elements] for layer in dopant_dict]
+        concentrations = np.array(vals)
+    elif concentrations is not None:
+        # Add Y into the list
+        if len(elements) != concentrations.shape[1]:
+            concentrations = np.concatenate(
+                (concentrations,
+                 1 - concentrations.sum(axis=1, keepdims=True)),
+                axis=1)
+
+    concentrations = np.clip(concentrations, 0, 1)
     colors = [
         DEFAULT_COLOR_MAP[el]
         if el in DEFAULT_COLOR_MAP else DEFAULT_COLOR_MAP['Other']
         for el in elements
     ]
-    # cmap = plt.colormaps["tab10"]
-    # colors = cmap(np.arange(4))
-    # # colors[:3] = colors[1:]
-    # colors[-1] = [1, 1, 1, 1]
-
-    fig = plt.figure(figsize=(5, 5), dpi=dpi)
-    ax = fig.subplots()
 
-    for i in range(concentrations.shape[0], 0, -1):
-        ax.pie(concentrations_with_y[i - 1],
-               radius=radii[i] / radii[-1],
-               colors=colors,
-               wedgeprops=dict(edgecolor='k', linewidth=0.25),
-               startangle=90)
-    ax.legend(elements, loc='upper left', bbox_to_anchor=(0.84, 0.95))
-    plt.tight_layout()
-    if as_np_array:
-        # If we haven't already shown or saved the plot, then we need to
-        # draw the figure first.
-        fig.canvas.draw()
-
-        # Now we can save it to a numpy array.
-        data = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
-        data = data.reshape(fig.canvas.get_width_height()[::-1] + (3, ))
-
-        # Close the figure to remove it from the buffer
-        plt.close(fig)
-        return data
+    if ax is None:
+        # make a new axis
+        fig = plt.figure(figsize=(5, 5), dpi=dpi)
+        ax = fig.subplots()
+
+        for i in range(concentrations.shape[0], 0, -1):
+            ax.pie(concentrations[i - 1],
+                   radius=radii[i] / radii[-1],
+                   colors=colors,
+                   wedgeprops=dict(edgecolor='w', linewidth=0.25),
+                   startangle=90)
+        ax.legend(elements, loc='upper left', bbox_to_anchor=(0.84, 0.95))
+        if emissions:
+            plt.text(0.1,
+                     0.95,
+                     f'UV Intensity={np.power(10, -emissions)-100:.2f}',
+                     fontsize=20,
+                     transform=plt.gca().transAxes)
+        plt.tight_layout()
+        if as_np_array:
+            # If we haven't already shown or saved the plot, then we need to
+            # draw the figure first.
+            fig.canvas.draw()
+
+            # Now we can save it to a numpy array.
+            data = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
+            data = data.reshape(fig.canvas.get_width_height()[::-1] + (3, ))
+
+            # Close the figure to remove it from the buffer
+            plt.close(fig)
+            return data
+        else:
+            return fig
     else:
-        return fig
-
-
-def plot_nanoparticle(constraints,
-                      dopant_specifications,
-                      dpi=150,
-                      as_np_array=False,
-                      elements=['Yb', 'Er', 'Nd']):
-    if 'Y' not in elements:
-        elements = elements + ['Y']
-
-    n_layers = len(constraints)
-    radii = [0] + [constraint.radius for constraint in constraints]
-    dopant_dict = [{key: 0 for key in elements} for _ in range(n_layers)]
-    for dopant in dopant_specifications:
-        dopant_dict[dopant[0]][dopant[2]] = dopant[1]
-
-    # Fill in the rest with 'Y'
-    for layer in dopant_dict:
-        layer['Y'] = 1 - sum(layer.values())
-
-    vals = [[layer[el] for el in elements] for layer in dopant_dict]
-
-    return plot_nanoparticle_from_arrays(np.array(radii),
-                                         np.array(vals),
-                                         dpi=dpi,
-                                         as_np_array=as_np_array,
-                                         elements=elements)
-
-
-def plot_nanoparticle_on_ax(ax,
-                            constraints,
-                            dopant_specifications,
-                            elements=['Yb', 'Er', 'Nd']):
-    if 'Y' not in elements:
-        elements = ['Y'] + elements
-
-    n_layers = len(constraints)
-    radii = [constraint.radius for constraint in constraints]
-    dopant_dict = [{key: 0 for key in elements} for _ in range(n_layers)]
-    for dopant in dopant_specifications:
-        dopant_dict[dopant[0]][dopant[2]] = dopant[1]
-    # Fill in the rest with 'Y'
-    for layer in dopant_dict:
-        layer['Y'] = np.round(1 - sum(layer.values()), 3)
-
-    vals = [[layer[el] for el in elements] for layer in dopant_dict]
-    cmap = plt.colormaps["tab10"]
-    colors = cmap(np.arange(4) * 4)
-    colors[0] = [1, 1, 1, 1]
-
-    for i in list(range(n_layers - 1, -1, -1)):
-        # print(vals[i])
-        ax.pie(vals[i],
-               radius=radii[i] / radii[-1],
-               colors=colors,
-               wedgeprops=dict(edgecolor='k'),
-               startangle=90)
-    ax.legend(elements, loc='upper left', bbox_to_anchor=(1, 1))
+        for i in range(concentrations.shape[0], 0, -1):
+            ax.pie(concentrations[i - 1],
+                   radius=radii[i] / radii[-1],
+                   colors=colors,
+                   wedgeprops=dict(edgecolor='w', linewidth=0.25),
+                   startangle=90)
+        ax.legend(elements, loc='upper left', bbox_to_anchor=(0.84, 0.95))
+        if emissions:
+            plt.text(0.1,
+                     0.95,
+                     f'UV Intensity={np.power(10, -emissions)-100:.2f}',
+                     fontsize=20,
+                     transform=plt.gca().transAxes)
 
 
 def update(data, ax):
-    constraints, dopants = data
     ax.clear()
-    plot_nanoparticle_on_ax(ax, constraints, dopants)
+    plot_nanoparticle(ax=ax, **data)
 
 
 def make_animation(frames: List[Tuple[NanoParticleConstraint, Tuple]],
                    name: str = 'animation.mp4',
-                   fps: int = 30) -> None:
+                   fps: int = 30,
+                   dpi: int = 300) -> None:
 
-    fig = plt.figure(dpi=150)
+    fig = plt.figure(dpi=dpi)
     ax = fig.subplots()
     anim = animation.FuncAnimation(fig, partial(update, ax=ax), frames=frames)
     anim.save(name, fps=fps)