diff --git a/src/NanoParticleTools/builder.py b/src/NanoParticleTools/builder.py index acec558..cb1bd04 100644 --- a/src/NanoParticleTools/builder.py +++ b/src/NanoParticleTools/builder.py @@ -53,12 +53,14 @@ def __init__(self, **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"]) diff --git a/src/NanoParticleTools/machine_learning/data/utils.py b/src/NanoParticleTools/machine_learning/data/utils.py index 5ab5aa8..d3c67c3 100644 --- a/src/NanoParticleTools/machine_learning/data/utils.py +++ b/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']) } diff --git a/src/NanoParticleTools/machine_learning/modules/ensemble.py b/src/NanoParticleTools/machine_learning/modules/ensemble.py index 3306c6f..5815b24 100644 --- a/src/NanoParticleTools/machine_learning/modules/ensemble.py +++ b/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) diff --git a/src/NanoParticleTools/optimization/callbacks.py b/src/NanoParticleTools/optimization/callbacks.py index 54b9de5..8c3e7bc 100644 --- a/src/NanoParticleTools/optimization/callbacks.py +++ b/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 diff --git a/src/NanoParticleTools/optimization/scipy_optimize.py b/src/NanoParticleTools/optimization/scipy_optimize.py index c802672..dab58d8 100644 --- a/src/NanoParticleTools/optimization/scipy_optimize.py +++ b/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 diff --git a/src/NanoParticleTools/species_data/species.py b/src/NanoParticleTools/species_data/species.py index d0169d7..bea6eee 100644 --- a/src/NanoParticleTools/species_data/species.py +++ b/src/NanoParticleTools/species_data/species.py @@ -84,6 +84,14 @@ class Dopant(MSONable): # 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 @@ class Dopant(MSONable): 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 diff --git a/src/NanoParticleTools/util/visualization.py b/src/NanoParticleTools/util/visualization.py index 946b67b..81c5f52 100644 --- a/src/NanoParticleTools/util/visualization.py +++ b/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) diff --git a/tests/util/test_visualization.py b/tests/util/test_visualization.py index 3464c98..07e3568 100644 --- a/tests/util/test_visualization.py +++ b/tests/util/test_visualization.py @@ -1,6 +1,5 @@ -from NanoParticleTools.util.visualization import ( - plot_nanoparticle_from_arrays, plot_nanoparticle, plot_nanoparticle_on_ax, - make_animation) +from NanoParticleTools.util.visualization import (plot_nanoparticle, + make_animation) from NanoParticleTools.inputs.nanoparticle import SphericalConstraint from matplotlib import pyplot as plt @@ -17,9 +16,9 @@ def test_plot_nanoparticle_from_arrays(): radii = np.array([0, 20, 50]) concentration = np.array([[0.25, 0.15, 0], [0.52, 0, 0.45]]) - plot_nanoparticle_from_arrays(radii, concentration) + plot_nanoparticle(radii, concentration) - out = plot_nanoparticle_from_arrays(radii, concentration, as_np_array=True) + out = plot_nanoparticle(radii, concentration, as_np_array=True) assert isinstance(out, np.ndarray) @@ -34,7 +33,7 @@ def test_plot_nanoparticle(): dopant_specifications = [(0, 0.25, 'Yb', 'Y'), (0, 0.15, 'Nd', 'Y'), (1, 0.52, 'Yb', 'Y'), (1, 0.45, 'Er', 'Y')] - plot_nanoparticle(constraints, dopant_specifications) + plot_nanoparticle(constraints, dopant_specifications=dopant_specifications) def test_plot_nanoparticle_on_ax(): @@ -50,7 +49,9 @@ def test_plot_nanoparticle_on_ax(): fig = plt.figure(dpi=100) ax = fig.add_subplot(121) - plot_nanoparticle_on_ax(ax, constraints, dopant_specifications) + plot_nanoparticle(constraints, + dopant_specifications=dopant_specifications, + ax=ax) def test_make_animation(): @@ -68,5 +69,8 @@ def test_make_animation(): (1, 0.52, 'Yb', 'Y'), (1, 0.45, 'Er', 'Y')], [(0, 0.25, 'Yb', 'Y'), (0, 0.15, 'Nd', 'Y'), (1, 0.52, 'Yb', 'Y'), (1, 0.35, 'Er', 'Y')]] - frames = list(zip(constraints, dopant_specifications)) + frames = [{ + 'radii': c, + 'dopant_specifications': d + } for c, d in zip(constraints, dopant_specifications)] make_animation(frames, name='animation.gif')