diff --git a/.gitignore b/.gitignore index 4f7f04f..26caaea 100644 --- a/.gitignore +++ b/.gitignore @@ -28,6 +28,7 @@ test_results/ !tests/outputs/figs/*.png *.jpg *.jpeg +*.sea # IDE and editor files .vscode/ diff --git a/src/pyslice/backend.py b/src/pyslice/backend.py index 61b1500..3d9a6a5 100644 --- a/src/pyslice/backend.py +++ b/src/pyslice/backend.py @@ -106,6 +106,8 @@ def asarray(arraylike, dtype=None, device=None): return array def zeros(dims, dtype=DEFAULT_FLOAT_DTYPE, device=DEFAULT_DEVICE): + if isinstance(dtype,str): + dtype=DEFAULT_FLOAT_DTYPE if dtype=="float" else DEFAULT_COMPLEX_DTYPE if xp != np: array = xp.zeros(dims, dtype=dtype, device=device) else: diff --git a/src/pyslice/data/signal.py b/src/pyslice/data/ALSONOTsignal.py similarity index 99% rename from src/pyslice/data/signal.py rename to src/pyslice/data/ALSONOTsignal.py index fac9a1b..272a383 100644 --- a/src/pyslice/data/signal.py +++ b/src/pyslice/data/ALSONOTsignal.py @@ -271,6 +271,7 @@ def to_dict(self, if deep: value = type(value)([v.to_dict(hidden=hidden, properties=properties, deep=deep) for v in value]) #keep the iterable type and loop it converting each element to a dict\ result[key] = value else: + #print("TODICT",key,value) result[key] = value return result @@ -282,6 +283,8 @@ def to_hdf5_group(self, parent_group: File|Group, """Save object to SEA formated HDF5 group.""" sea_type = type(self).__name__ # Get the SEA class type + if sea_type in ["TACAWData","WFData"]: # TWP 20260120: sea-eco appears to not know about TACAWData, so let's treat parent classes (not Dimensions and friends) as Signals + sea_type = 'Signal' # Handle group name: use provided name, fall back to object's name attribute, or default if name is None: @@ -299,8 +302,16 @@ def to_hdf5_group(self, parent_group: File|Group, exclude_keys=exclude_keys) to_write = {**{k:v for k,v in to_write.items() if '_' not in str(k)}, **{k:v for k,v in to_write.items() if '_' in str(k)}} + print("to_write.keys()",to_write.keys()) + #to_write[ "data" ]=to_write[ "array" ] storage_name_counts: Dict[str, int] = {} + #print("to_write",to_write.items().keys()) for key, val in to_write.items(): + #print("name",name,"key",key,"val",val,"type",type(val)) + if key in ["probe","cache_dir"]: + continue + if key == "_array": + key = "data" if not hasattr(self, key): continue storage_key = key[1:] if key.startswith('_') else key @@ -339,7 +350,7 @@ def to_sea(self, file_path: str, file.attrs['file_type'] = 'SEA-eco HDF5 file'.encode('utf-8') file.attrs['file_version'] = '0.0'.encode('utf-8') file.attrs['sea_type'] = type(self).__name__.encode('utf-8') - + #print("file.attrs",[ (k,file.attrs[k]) for k in file.attrs.keys() ]) self.to_hdf5_group(parent_group=file, force_datasets=force_datasets) file.close() @@ -632,7 +643,7 @@ class Dimension(SEASerializable): def __init__(self, dimension: Dict|GeneralMetadata|None = None, name: str = 'Unnamed Dimension', space: Literal["position", "scattering", "temporal", "spectral"] | None = None, - scale: float|int|Iterable[float|int]|None = None, + scale: float|int|Iterable[float|int]|None = 1, offset: float|int|Iterable[float|int]|None = None, size: int|Iterable[int]|None = None, units: str|Iterable[str] = '', diff --git a/src/pyslice/data/NOTsignal.py b/src/pyslice/data/NOTsignal.py new file mode 100644 index 0000000..e627826 --- /dev/null +++ b/src/pyslice/data/NOTsignal.py @@ -0,0 +1,1802 @@ +''' +Signal class and related data structures for PySlice. +Copied from sea-eco with minimal modifications. +''' + +#Imports: Typing +from __future__ import annotations +from abc import ABC, abstractmethod +from typing import List, Any, Dict, Tuple, Callable, Literal, Union +#from types import EllipsisType +from collections.abc import Iterable +from matplotlib.axes import Axes as mplAxes +import matplotlib.pyplot as plt +from numpy.typing import NDArray, DTypeLike + +#Imports: External +from warnings import warn +from copy import deepcopy +from pathlib import Path +from inspect import signature, Parameter + +from h5py import Group, File, Dataset +from uuid import uuid4 + +import numpy as np +import pickle + +# Optional plotting imports from sea-eco +try: + from pySEA.sea_eco._plotting.plot import plot_nd_array, PlotImage, save_plot, save_image + PLOTTING_AVAILABLE = True +except ImportError: + PLOTTING_AVAILABLE = False + plot_nd_array = None + PlotImage = None + save_plot = None + save_image = None + +def generate_uuid() -> str: + """Generate a new UUID string.""" + return str(uuid4()) + +def _check_and_convert_numpy(value: Any) -> Any: + """Convert numpy scalars/arrays to native Python containers.""" + if isinstance(value, np.generic): + return _check_and_convert_numpy(value.item()) + #if isinstance(value, np.ndarray): + # return [_check_and_convert_numpy(v) for v in value.tolist()] + if isinstance(value, (list, tuple)): + converted = [_check_and_convert_numpy(v) for v in value] + return type(value)(converted) + return value + +def safe_decode(value: Any) -> Any: + """Safely decode bytes and numpy values to native Python objects. + + Used for HDF5 attributes (not datasets). Converts bytes to strings + and handles special encoded values like None. + """ + if isinstance(value, bytes): + try: + decoded = value.decode('utf-8') + # Handle special encoded values + if decoded == 'None': + return None + return decoded + except UnicodeDecodeError: + return value + if isinstance(value, str): + # Handle string 'None' (in case it wasn't bytes) + if value == 'None': + return None + return _check_and_convert_numpy(value) + +def safe_encode(value: Any) -> Any: + """Safely encode Python values (including numpy types) for HDF5 attributes.""" + if isinstance(value, str): + return value.encode('utf-8') + if value is None: + return 'None'.encode('utf-8') + return _check_and_convert_numpy(value) + +def ask_to_proceed(): + while True: + user_input = input("Do you want to continue? (Y/N): ").upper() # Convert to uppercase for case-insensitivity + if user_input == 'Y' or user_input== '': + print("Proceeding...") + # Add the code to execute if the user chooses 'Y' here + break # Exit the loop once a valid 'Y' is entered + elif user_input == 'N': + print("Exiting...") + # Add the code to execute if the user chooses 'N' here + break # Exit the loop once a valid 'N' is entered + else: + print("Invalid input. Please enter 'Y' or 'N'.") + +def get_index_or_none(list_in:List, value:Any): + try: + index = list_in.index(value) + return index + except ValueError: + return None + +def check_dimensions_call(kwargs, fnc): + """Check and convert dimension kwargs (dim, axis, axes) to the function's expected format then returned the allowed kwarg key. + + Parameters + ---------- + kwargs : dict + The keyword arguments to check and potentially modify + fnc : callable + The function whose signature to check + + Returns + ------- + str | None + The dimension keyword that was used, or None if no dimension kwargs were found + """ + possible_dim_keys = ['axis', 'axes', 'dim'] + + # First check what dimension arguments the function accepts + params = signature(fnc).parameters + fnc_accepts = None + for k in possible_dim_keys: + if k in params: + fnc_accepts = k + break + if fnc_accepts is None: + warn('The function does not accept any dimensional references.') + + # Check which dimension arguments were provided + provided_key = [k for k in possible_dim_keys if k in kwargs] + + if len(provided_key) == 0: + return None + elif len(provided_key) == 1: + kwargs[fnc_accepts] = kwargs.pop(provided_key[0]) + return fnc_accepts + elif len(provided_key) > 1: + raise ValueError(f"Only one of {possible_dim_keys} can be used at a time, but got: {provided_key}") + + #kwargs[fnc_accepts] = kwargs.pop(provided_key) + + return fnc_accepts #provided_key #? should I provide the one that was found in fnc or provided + + # if prov_dim_key: kwargs[fnc_dim_key] = kwargs.pop(prov_dim_key) + # return prov_dim_key + +def get_property_dict(obj:object): + """Get a dictionary of an objects properties + + Parameters + ---------- + obj : Object + Any object + + Returns + ------- + Dict + A dictionary of an objects properties. + """ + cls = type(obj) + return { + name: getattr(obj, name) + for name in dir(cls) + if isinstance(getattr(cls, name), property) + } + +def get_tree_html(obj, recursive_level: int = 0, + exclude_keys: List[str] = [], + exclude_hidden: bool = True, + exclude_properties: bool = False, + promote_itterable_keys: List[str] = [] + ) -> str: + html = "" + tree_dict = dict(obj.__dict__) + if not exclude_properties: tree_dict.update(get_property_dict(obj)) + + for key, value in tree_dict.items(): + #Handle exclusions and promotions of keys + if key in exclude_keys: continue + if exclude_hidden and key.startswith('_'): continue + + #build the html tree + if isinstance(value, SEASerializable): + html += ( + "" + ) + elif isinstance(value, Iterable) and key in promote_itterable_keys: + html += ( + "" + else: + html += ( + "" + ) + return html + +class SEASerializable(ABC): + """Base class for objects that can be serialized to/from HDF5.""" + + def _repr_html_(self): + """Return a compact HTML representation of the GeneralMetadata tree.""" + return self.get_tree_html() + + def deepcopy(self): + """Create a deep copy of the object.""" + return deepcopy(self) + + def to_dict(self, + hidden: bool = True, + properties: bool = False, + exclude_keys: List[str] = [], + deep: bool = True + #promote_itterable_keys: List[str] = [] # Iterables atributes like Dimensions.dimensions or Signals.signals to promote. #?This was used in get_tree_html but I am not sure if it will be needed here. + ): + """ + Recursively convert Dimension object to a dictionary. + + Parameters + ---------- + hidden : bool, optional + Include hidden attributes (starting with '_'), by default Truee. + properties : bool, optional + Include properties, by default False. + exclude_keys : List[str], optional + Keys to exclude from the dictionary, by default []. + deep : bool, optional + Recursively convert nested SEASerializable objects, by default True. + + Returns + ------- + dict + Dictionary representation of the Dimension object. + """ + to_convert = dict(self.__dict__) + if properties: to_convert.update(get_property_dict(self)) + + result = {} + for key, value in to_convert.items(): + if not hidden and key.startswith('_'): continue #if converting hiddens then continue to the next key + if key in exclude_keys: continue #if an exclueded key then continue to the next key + + if isinstance(value, SEASerializable): + if deep: value = value.to_dict(hidden=hidden, properties=properties, deep=deep) + result[key] = value + elif isinstance(value, dict): # In case a dict is stored directly + if deep: value = {k: v.to_dict(deep=deep) if isinstance(v, SEASerializable) else v for k, v in value.items()} + result[key] = value + elif isinstance(value, Tuple | List) and len(value) > 0 and isinstance(value[0], SEASerializable): #If an iterable of SEASerializable is given + if deep: value = type(value)([v.to_dict(hidden=hidden, properties=properties, deep=deep) for v in value]) #keep the iterable type and loop it converting each element to a dict\ + result[key] = value + else: + #print("TODICT",key,value) + result[key] = value + return result + + def to_hdf5_group(self, parent_group: File|Group, + force_datasets: List = [], + name: str | None = None, + exclude_keys: List[str] = [] + ) -> None: + """Save object to SEA formated HDF5 group.""" + + sea_type = type(self).__name__ # Get the SEA class type + if sea_type in ["TACAWData","WFData"]: # TWP 20260120: sea-eco appears to not know about TACAWData, so let's treat parent classes (not Dimensions and friends) as Signals + sea_type = 'Signal' + + # Handle group name: use provided name, fall back to object's name attribute, or default + if name is None: + if hasattr(self, 'name') and getattr(self, 'name') is not None: name = getattr(self, 'name') + else: name = sea_type + + # Create the group and assign SEA class type + group = parent_group.create_group(name, track_order=True) + group.attrs['sea_type'] = sea_type + + # Iteratively assign attributes by iterating to_dict. + # Keep the to_dict local so that to_hdf5 can be called for any serializable value. + to_write = self.to_dict(deep=False, hidden=True, + properties=False, + exclude_keys=exclude_keys) + to_write = {**{k:v for k,v in to_write.items() if '_' not in str(k)}, + **{k:v for k,v in to_write.items() if '_' in str(k)}} + print("to_write.keys()",to_write.keys()) + #to_write[ "data" ]=to_write[ "array" ] + storage_name_counts: Dict[str, int] = {} + #print("to_write",to_write.items().keys()) + for key, val in to_write.items(): + #print("name",name,"key",key,"val",val,"type",type(val)) + if key in ["probe","cache_dir"]: + continue + if key == "_array": + key = "data" + if not hasattr(self, key): + continue + storage_key = key[1:] if key.startswith('_') else key + storage_name_counts[storage_key] = storage_name_counts.get(storage_key, 0) + 1 + if storage_name_counts[storage_key] > 1: + warn(f'HDF5 serialization collision for attribute name \"{storage_key}\". Hidden and public attributes share the same name.') + #Check the value to see if it is Serializable + if isinstance(val, SEASerializable): + val.to_hdf5_group(parent_group=group, name=storage_key) + continue + elif isinstance(val, Iterable) and not isinstance(val, str) and len(val) > 0 and all(isinstance(v, SEASerializable) for v in val): + if key in force_datasets or storage_key in force_datasets: + key_group = group + else: + key_group = group.create_group(name=storage_key, track_order=True) + key_group.attrs['sea_type'] = val.__class__.__name__ #get the iterable type as string + for v in val: + v.to_hdf5_group(parent_group=key_group) + continue + + # Check the keys to see if it should be a attribute or dataset + if key in force_datasets or storage_key in force_datasets: + group.create_dataset(storage_key, data=val) + else: group.attrs[storage_key] = safe_encode(val) + return group + + def to_sea(self, file_path: str, + force_datasets: List = []) -> None: + """Save object to SEA formated HDF5.""" + file_path = Path(file_path) + if file_path.suffix != '.sea' and file_path.suffix != '': + raise ValueError("The file extension must be '.sea' or empty.") + file_path = str(file_path.with_suffix(''))+'.sea' + + file = File(file_path, "w") + file.attrs['file_type'] = 'SEA-eco HDF5 file'.encode('utf-8') + file.attrs['file_version'] = '0.0'.encode('utf-8') + file.attrs['sea_type'] = type(self).__name__.encode('utf-8') + #print("file.attrs",[ (k,file.attrs[k]) for k in file.attrs.keys() ]) + self.to_hdf5_group(parent_group=file, force_datasets=force_datasets) + + file.close() + + def from_sea(self, file_path: str): + """Load object data from an HDF5 file.""" + file_path = Path(file_path) + if file_path.suffix != '.sea' and file_path.suffix != '': + raise ValueError("The file extension must be '.sea' or empty.") + file_path = str(file_path.with_suffix(''))+'.sea' + file = File(file_path, "r") + + if len(file)!=1: + raise ValueError("The hdf5 file contains multiple groups so can not be loaded directly. Consider using `from_hdf5_group` instead to append to the current class.") + else: + main_group = file[list(file.keys())[0]] + + if 'sea_type' not in main_group.attrs: + raise ValueError("Could not locate an HDF5 group matching this object.") + elif safe_decode(main_group.attrs['sea_type']) != type(self).__name__: + raise ValueError(f"The HDF5 group sea_type '{safe_decode(main_group.attrs['sea_type'])}' does not match the current object type '{type(self).__name__}'.") + else: + self.from_hdf5_group(main_group) + + file.close() + + def from_hdf5_group(self, group: Group): + """Populate the current object from an HDF5 group. + """ + def _instantiate_child(sub_group: Group): + """Instantiate an SEASerializable child using a prototype or group metadata.""" + candidate = None + sea_type = safe_decode(sub_group.attrs.get('sea_type', b'')) + cls = globals().get(sea_type) + if isinstance(cls, type) and issubclass(cls, SEASerializable): + try: + candidate = cls() + except TypeError: + candidate = cls.__new__(cls) + cls.__init__(candidate) + if candidate is None: return None + candidate.from_hdf5_group(sub_group) + return candidate + def _check_attr_visibility(attr: str) -> str | None: + """_check_attr_visibility _summary_ + + Parameters + ---------- + attr : str + Check if the attribute exists as public or private and return the correct name. + + Returns + ------- + str | None + _description_ + """ + if hasattr(self, f'_{attr}'): return f'_{attr}' + else: return attr + + for key, val in group.attrs.items(): + if key == 'sea_type': continue + attr_name = _check_attr_visibility(key) + if attr_name is not None: + setattr(self, attr_name, safe_decode(val)) + + for key, item in group.items(): + if isinstance(item, Group): + if 'sea_type' in item.attrs: + sea_type = safe_decode(item.attrs.get('sea_type', b'')) + else: + warn(f'Group {item.name} has no sea_type attribute. Skipping.') + continue + attr_name = _check_attr_visibility(key) + if attr_name is None: + warn(f'Attribute {key} from group {item.name} not found on target object.') + continue + if sea_type=='list': + new_items = [] + for sub_item in item.values(): + child = _instantiate_child(sub_item) + if child is not None: + new_items.append(child) + setattr(self, attr_name, new_items) + else: + child = _instantiate_child(item) + if child is not None: + setattr(self, attr_name, child) + else: + warn(f'Attribute {key} could not be instantiated from group {item.name}. Skipping.') + else: + # This is an HDF5 dataset (array data) + value = item[()] + # For datasets, preserve numpy arrays - only decode bytes/strings + if isinstance(value, bytes): + value = safe_decode(value) + elif isinstance(value, np.ndarray) and value.dtype.kind in ('S', 'U', 'O'): + # String/object arrays - decode elements + value = np.array([safe_decode(v) for v in value.flat]).reshape(value.shape) + # Otherwise keep as numpy array + attr_name = _check_attr_visibility(key) + if attr_name is None: + warn(f'Dataset {key} could not be assigned to target object.') + continue + setattr(self, attr_name, value) + + def get_tree_html(self, recursive_level:int=0, + exclude_keys:List[str] = [], + exclude_hidden:bool=True, + exclude_properties:bool=False, + promote_itterable_keys: List[str] = [] + ) -> str: + #TODO: Have this work of of self.to_dict() instead of self.__dict__ and put the global get_tree_html() kwargs in self.to_dict + return get_tree_html(self, recursive_level=recursive_level, + exclude_keys=exclude_keys, + exclude_hidden=exclude_hidden, + exclude_properties=exclude_properties, + promote_itterable_keys=promote_itterable_keys) + + def get_tree_str(self, pad:str='', recursive_level=None): + """Print class as a tree. + + Parameters + ---------- + pad : str, optional + What to add before each entry, by default '' + recursive_level : int, optional + What depth to stop recursing. Not implemented, by default 0 + + Returns + ------- + _type_ + _description_ + """ + if recursive_level is not None: + #TODO: have the tree stop at a level + raise NotImplementedError('kwarg recursive_level not implemented yet.') + string = '' + N_values = len(self) + for i, (key, value) in enumerate(self.__dict__.items()): + if i==N_values-1: cnct = '└── ' + else: cnct = '├── ' + + if isinstance(value, GeneralMetadata): + string += f'{pad}{cnct}{key}\n' + if i== N_values-1: pad_next = ' ' + else: pad_next = '| ' + string += value.get_tree_str(pad=pad+pad_next) + else: + string += f'{pad}{cnct}{key}: {value}\n' + return string + + def show_tree(self, show:Literal['html','str']|None='html', recursive_level=0): + if show=='html': + from IPython.display import display, HTML + display(HTML(self.get_tree_html(recursive_level=recursive_level))) + elif show=='str': + print(self.get_tree_str(self, recursive_level=recursive_level)) + else: + try: + from IPython.display import display, HTML + display(HTML(self.get_tree_html(recursive_level=recursive_level))) + except: + print(self.get_tree_str(self, recursive_level=recursive_level)) + + def save(self, file_path: str) -> None: + file_path = Path(file_path) + if file_path.suffix == '.sea' or file_path.suffix == '': + self.to_sea(file_path) + elif file_path.suffix == '.pkl': + with open(file_path,'wb') as f: pickle.dump(self,f) + else: + raise ValueError("The file extension must be '.sea', '.pkl' or empty.") + +def load(file_path: str) -> SEASerializable: + """Load a SEASerializable object from a .sea or .pkl file. + + Parameters + ---------- + file_path : str + Path to the file to load. + + Returns + ------- + SEASerializable + The loaded object. + """ + file_path = Path(file_path) + + if file_path.suffix == '.sea' or file_path.suffix == '': + file_path = str(file_path.with_suffix('')) + '.sea' + with File(file_path, 'r') as f: + if len(f) != 1: + raise ValueError("The HDF5 file contains multiple groups. Cannot determine which to load.") + main_group = f[list(f.keys())[0]] + sea_type = safe_decode(main_group.attrs.get('sea_type', b'')) + + # Get the class from globals and instantiate + cls = globals().get(sea_type) + if cls is None: + raise ValueError(f"Unknown sea_type: {sea_type}") + + # Create instance and load data + try: + obj = cls() + except TypeError: + obj = cls.__new__(cls) + obj.from_sea(file_path) + return obj + elif file_path.suffix == '.pkl': + with open(file_path, 'rb') as f: + obj = pickle.load(f) + return obj + +class GeneralMetadata(SEASerializable): + def __init__(self, meta:dict = {}) -> None: + """General Metadata + + Parameters + ---------- + meta : dict + Dictionary to be converted to metadata. + """ + self.update_from_dict(dictionary=meta) + + def __len__(self): + """Return the number of items in the GeneralMetadata.""" + return len(self.__dict__) + + def __repr__(self): + """Return a compact HTML representation of the GeneralMetadata tree.""" + return self.get_tree_str() + + def to_hdf5_group(self, parent_group, force_datasets = [], name = 'Metadata'): + return super().to_hdf5_group(parent_group, force_datasets, name) + + def update_from_dict(self, dictionary:Dict[str, Dict | List | Any | None]) -> None: + '''Recursively update GeneralMetadata object from a dictionary. + + Parameters + ---------- + dictionary : dict + The key will define the Node name. + If the value is an empty dictionary the value will be asigned as None. + If the value is a dictionary then the value will be another GeneralMetadata object. + Otherwise the value will be asigned directly.''' + for key, value in dictionary.items(): + if isinstance(value, dict): + if len(value)==0: setattr(self, key, None) + else: setattr(self, key, GeneralMetadata(value)) + else: + setattr(self, key, value) + + def merge(self, other: 'GeneralMetadata', + kind:Literal['overwrite','append','skip']='skip', warn_duplicate:bool=False, + inform_new=False) -> None: + """Merge another GeneralMetadata object into this one. + + Parameters + ---------- + other : GeneralMetadata + The other GeneralMetadata object to merge. + kind : bool, optional + If True, existing keys will be overwritten. Default is False. + warn : bool, optional + If True, warn when duplicates arise. Default is False. + """ + for key, value in other.__dict__.items(): + if hasattr(self, key): + if isinstance(getattr(self, key), GeneralMetadata) and isinstance(value, GeneralMetadata): + getattr(self, key).merge(value, kind=kind, warn_duplicate=warn_duplicate, inform_new=inform_new) + elif kind=='overwrite': + setattr(self, key, value) + if warn_duplicate: warn(f'Duplicate - Key {key} already exists. Overwriting.') + elif kind=='append': + existing = getattr(self, key) + if existing != value: + if not isinstance(value, List): value = [value] + if not isinstance(existing, List): existing = [existing] + setattr(self, key, value+existing) + if warn_duplicate: warn(f'Duplicate - Key {key} already exists. Appending.') + elif warn_duplicate: warn(f'Duplicate - Key {key} already exists. Use kind=True to overwrite.') + else: + setattr(self, key, value) + if inform_new: print(f'New - Key {key} added to metadata.') + +class Dimension(SEASerializable): + """ + Dimensions of a dataset. + """ + def __init__(self, dimension: Dict|GeneralMetadata|None = None, + name: str = 'Unnamed Dimension', + space: Literal["position", "scattering", "temporal", "spectral"] | None = None, + scale: float|int|Iterable[float|int]|None = 1, + offset: float|int|Iterable[float|int]|None = None, + size: int|Iterable[int]|None = None, + units: str|Iterable[str] = '', + values: NDArray[Any] | Iterable[Any] = None, + #device: torch.device = field(default_factory=get_default_device) + ) -> None: + """Initialize an Dimension object from a dictionary or GeneralMetadata object. + + Parameters + ---------- + dimension : Dict | GeneralMetadata | None, optional + Dict or metadata object containing, by default None + name : str, optional + Name of the dimension, by default 'Unnamed Dimension' + space : Literal["position", "scattering", "temporal", "spectral"] | None, optional + Space the dimension belongs to, by default None + scale : float | int | Iterable[float | int] | None, optional + Calibration scale, by default None + offset : float | int | Iterable[float | int] | None, optional + Calibration offset, by default None + size : int | Iterable[int] | None, optional + Dimension size, by default None + units : str | Iterable[str], optional + Calibration units, by default '' + aixs : NDArray[Any] | Iterable[Any], optional + Values along the dimension, by default None + + ToDo + ---- + TODO: Add a `parametric:bool=True` and `parametric_fnc:str='x*scale+offset'` attributes + """ + #Initialize instance attributes + self.dimension = dimension + self._name = name + self.space = space + self.scale = scale + self.offset = offset + self.size = size + self.units = units + self.values = values + + defaults = {'name':'Unnamed Dimension', 'space':None, 'scale':None, 'offset':None, 'units':'', 'size':None} #self.__init__.__kwdefaults__ + if self.dimension is not None and not isinstance(self._values, Iterable): #if an dimension and not values is given + if isinstance(dimension, GeneralMetadata): self.dimension = self.dimension.to_dict() #convert to dict if GeneralMatadata + for k,v in self.dimension.items(): #loop dimension dict and + if k in defaults.keys() and defaults[k]==getattr(self,k): #asign the key if it is equal to the default + setattr(self, k, v) + elif self.dimension is None and isinstance(self._values, Iterable): #if values and not an dimension is given + self.size = self._values.shape[-1] #HACK: this should not be size but some sort of tuple that acounts for ndim. + elif self.dimension is None and not isinstance(self._values, Iterable): #if no values or dimensions are supplied we set the offset and scale. We don't do it by default because these should be None if the dimension is not parametric. + if self.offset is None: self.offset = 0 + if self.scale is None: self.scale = 1 + else: + self._values = None + del(self.dimension) + + def __str__(self): + return f'{self.name}-dimension' + + def __repr__(self): + return f'' + + def __getitem__(self, key: Union[int, float, slice, tuple]): + """Support indexing and slicing of the Signal. + + Parameters + ---------- + key : Union[int, float, slice, tuple, EllipsisType] + Index specification. Can include integers, floats (converted to nearest index), + slices, ellipsis, or tuples of these. + - int: direct indexing + - float: converted to nearest index using dimension calibration + - slice: regular Python slicing, can include float values + - Ellipsis: expands to cover remaining dimensions + """ + if not isinstance(key, tuple): key = (key,) + + key_reg = tuple() + for i, k in enumerate(key): + if isinstance(k, float): + k = self.find_nearest_index(k) + elif isinstance(k, slice): + start = k.start + stop = k.stop + step = k.step + if isinstance(start, float): start = self.find_nearest_index(start) + if isinstance(stop, float): stop = self.find_nearest_index(stop) + if isinstance(step, float): step = int(round(step / self.scale, 0)) + k = slice(start, stop, step) + else: + pass # Ellipse or int do not remove + key_reg += (k,) + + return self.values[key_reg] + + def _check_ndim(self): + checks = {'size':self.size, + 'offset':self.offset, + 'scale':self.scale} + if all((var is not None for var in checks.values())): + n_scale = len(self.scale) if isinstance(self.scale,Iterable) else 1 + n_offset = len(self.offset) if isinstance(self.offset,Iterable) else 1 + n_units = len(self.units) if isinstance(self.units,Iterable) and not isinstance(self.units,str) else 1 + if n_scale>1 and n_offset>1 and n_scale!=n_offset: + raise ValueError('Scale and offset are >1 dimensions but not consistent dimensionality.') + if n_units>2 and n_units!=max(n_scale,n_offset): + raise ValueError('Units is >1 dimensions but not consistent dimensionality with the larger of scale or offset.') + return max(n_scale,n_offset) + elif self._values is not None: + print(self._values) + return np.ndim(self._values) + else: + warn('The dimensions chould not be determined from the calibrations or values.') + return None + + @property + def ndim(self): + return np.ndim(self.values) + @ndim.setter + def ndim(self, values: Any) -> UserWarning: + #if values is None: self._original_metadata = values + raise UserWarning('`ndim` is read-only.') + + @property + def values(self): + #if self._values is None: + checks = {'size':self.size, + 'offset':self.offset, + 'scale':self.scale} + if all((var is not None for var in checks.values())): + self._values = np.arange(self.size)*np.expand_dims(self.scale, axis=-1) + np.expand_dims(self.offset, axis=-1) + elif self._values is not None: + pass + else: + if all((var is None for var in checks.values())) and self._values is None: + pass + else: + for key, var in checks.items(): + if key not in checks.items(): warn(f'{key} is None.') + return self._values + @values.setter + def values(self, values): + self._values = values + + @property + def name(self): + if not isinstance(self._name, str) and isinstance(self._name, Iterable): + return '('+', '.join(self._name)+')' + else: return self._name + @name.setter + def name(self, value): + self._name = value + #self._check_ndim() + + def to_dict(self, + hidden: bool = False, + properties: bool = True, + exclude_keys: List[str] = [], + deep: bool = True + ): + """ + Recursively convert Dimension object to a dictionary. + + Parameters + ---------- + hidden : bool, optional + Include hidden attributes (starting with '_'), by default False. + properties : bool, optional + Include properties, by default True. + exclude_keys : List[str], optional + Keys to exclude from the dictionary, by default []. + deep : bool, optional + Recursively convert nested SEASerializable objects, by default True. + + Returns + ------- + dict + Dictionary representation of the Dimension object. + """ + return super().to_dict(hidden=hidden, properties=properties, exclude_keys=exclude_keys, deep=deep) + + def to_hdf5_group(self, parent_group: File|Group, + force_datasets: List = ['_values'], + name: str | None = None + ) -> None: + super().to_hdf5_group(parent_group=parent_group, + force_datasets=force_datasets, + name=name) + def to_sea(self, file_path: str, + force_datasets: List = ['_values']) -> None: + """Save dimension to HDF5.""" + super().to_sea(file_path, + force_datasets=force_datasets) + + def _get_tree_html(self, recursive_level: List[str] = 0, + exclude_keys: List[str]= ['values'], + exclude_hidden: bool = True, + exclude_properties:bool = False, + promote_itterable_keys: List[str] = [] + ) -> str: + return super().get_tree_html(recursive_level, + exclude_keys=exclude_keys, + exclude_hidden=exclude_hidden, + exclude_properties=exclude_properties, + promote_itterable_keys=promote_itterable_keys + ) + def get_tree_html(self, recursive_level: int = 0, + exclude_keys: List[str] = [], + exclude_hidden: bool = True, + exclude_properties:bool = False, + promote_itterable_keys: List[str] = [] + ) -> str: + return self._get_tree_html(recursive_level, + exclude_keys=exclude_keys+['values'], + exclude_hidden=exclude_hidden, + exclude_properties=exclude_properties, + promote_itterable_keys=promote_itterable_keys + ) + + def get_calibrated_value(self, indices: int | Iterable[int]) -> float: + """Get calibrated value at a specific index. + + Parameters + ---------- + indices : int | Iterable[int] + Indices in axis array. + + Returns + ------- + float + Value at the specified index. + + Raises + ------ + IndexError + Out of range index. + """ + if self.ndim > 1: + if len(indices) != self.ndim: raise ValueError(f"Expected {self.ndim} indices, got {len(indices)}") + for ind, dim in zip(indices, self.size): + if ind >= dim: raise IndexError(f"Index {ind} out of range for size {dim}") + else: + if indices >= self.size: raise IndexError(f"Index {indices} out of range for size {self.size}") + return float(self.values[indices]) + + def find_nearest_index(self, value: float, + direction:Literal['boht','above','below']='both', + warn_bounds=True) -> int: + """Find the index of the nearest calibrated value. + + Parameters + ---------- + value : float + Calibrated value to find the nearest index to. + direction : Literal['both','above','below'] + Direction to resolve ties. 'both' returns nearest, 'above' returns next higher index, 'below' returns next lower index. Default is 'both'. + warn_bounds : bool, optional + Warn if the nearest index is at the bounds of the axis, by default True. + + Returns + ------- + int + The nearest index. + """ + if self.ndim > 1: raise NotImplementedError('find_nearest_index only implemented for 1D dimensions.') #TODO: Implement for multi-D dimensions. + distances = np.abs(self.values - value) + index = int(np.argmin(distances)) + if warn_bounds and (index==0 or index==self.size-1): + print(f'Warning: Nearest index {index} is at the bounds of the axis (0, {self.size-1}).') + if direction == 'both': + return index + elif direction == 'above': + return index + 1 + elif direction == 'below': + return index - 1 + + def get_extent(self) -> List[float]| List[Tuple[float]]: + """Get the extent of the dimension for plotting. + + Returns + ------- + List[float] + Extent as [min, max]. + """ + if self.ndim == 1: + return [np.min(self.values[0]), np.max(self.values[-1])] + if self.ndim > 1: + return list(zip(np.min(self.values, axis=0), np.max(self.values, axis=0))) + +class Dimensions(SEASerializable): + def __init__(self, + dimensions: Iterable[Dict|Dimension] = [], + nav_dimensions: List[int] = [], + sig_dimensions: List[int] = [], + ) -> None: + self.dimensions = dimensions #! This could be set as private then have __get_item__ iterate the private class. + self.nav_dimensions = nav_dimensions + self.sig_dimensions = sig_dimensions + self.order = list(range(len(self.dimensions))) + #? The below hidden dimensions are storgage for the get/set property. The getter always stores in the hidden, so is there any point in having the hiddens? I guess as is, the set would store in the hidden, then that hidden could be accessed directly allowing the user to hack if needed. + self._spectral_dimension = [] + self._temporal_dimension = [] + self._position_dimensions = [] + self._scattering_dimensions = [] + + dimensions_list = [] + for i, dimension in enumerate(self.dimensions): + if isinstance(dimension, Dict): + dimension_obj = Dimension(dimension) + elif isinstance(dimension, Dimension): + dimension_obj = dimension + else: + raise TypeError(f'Dimensions iterable value of {type(dimension)} was provided but is not an allowed type.') + dimensions_list.append(dimension_obj) + self.dimensions = dimensions_list + + @property + def ndim(self) -> int: + """Get the total number of dimensions across all dimensions.""" + ndim = np.sum([ax.ndim for ax in self.dimensions], dtype=int) + return ndim + @ndim.setter + def ndim(self, value:Any) -> UserWarning: + raise UserWarning('ndim should not be set by the user.') + + @property + def spectral_dimension(self) -> int: + for i, dimension in enumerate(self.dimensions): + if dimension.space=='spectral': + self._spectral_dimension = i + break + return self._spectral_dimension + @spectral_dimension.setter + def spectral_dimension(self, value:int) -> None: + self._spectral_dimension = value + + @property + def temporal_dimension(self) -> int: + for i, dimension in enumerate(self.dimensions): + if dimension.space=='temporal': + self._temporal_dimension = i + break + return self._temporal_dimension + @temporal_dimension.setter + def temporal_dimension(self, value:int) -> None: + self._temporal_dimension = value + + @property + def position_dimensions(self) -> List: + if len(self._position_dimensions)==0: # first run only, assemble based on dimension.space + for i, dimension in enumerate(self.dimensions): + if dimension.space=='position': + self._position_dimensions.append(i) + return self._position_dimensions + @position_dimensions.setter + def position_dimensions(self, value:List) -> None: + self._position_dimensions = value + + @property + def scattering_dimensions(self) -> List: + if len(self._scattering_dimensions)==0: # first run only, assemble based on dimension.space + for i, dimension in enumerate(self.dimensions): + if dimension.space=='scattering': + self._scattering_dimensions.append(i) + return self._scattering_dimensions + @scattering_dimensions.setter + def scattering_dimensions(self, value:List) -> None: + self._scattering_dimensions = value + + def __repr__(self): + return f'' + + def __getitem__(self, key:int|str|Iterable[int|str]): + if isinstance(key, int): return self.dimensions[key] + elif isinstance(key, str): return self.dimensions[self.get_index_from_name(key)] + elif isinstance(key, Iterable): + ret = [] + for k in key: + if isinstance(k, int): ret.append(self.dimensions[k]) + elif isinstance(k, str): ret.append(self.dimensions[self.get_index_from_name(k)]) + return ret + else: raise TypeError(f'Only integers and strings are allowed but a {type(key)} was provided.') + + def __len__(self): + return len(self.dimensions) + + def add_dimension(self, dimension:Dict|Dimension) -> None: + """Add an dimension to the Dimensions object. + + Parameters + ---------- + dimension : Dict | Dimension + Dictionry with dimension calibrations or Dimension class to add. + + Raises + ------ + TypeError + If the provided dimension is not a dictionary or Dimension object. + """ + dimension_n = len(self.dimensions) + if isinstance(dimension, Dict): + dimension = Dimension(dimension) + if isinstance(dimension, Dimension): + self.dimensions.append(dimension) + else: raise TypeError(f'Dimensions iterable value of {type(dimension)} was provided but is not an allowed type.') + + self.order.append(dimension_n) + + def get_names(self) -> List[str]: + return [ax.name for ax in self.dimensions] + + def get_index_from_name(self, name:str) -> int: + names = self.get_names() + if name in names: return names.index(name) + else: raise KeyError(f'A key of {name} was provided and the dimensions names are {names}') + + def get_dims_as_int(self, + dims: str | int | Iterable[str | int] | None + ) -> List | int: + """Convert named or integer indices to integers. + + Parameters + ---------- + dims : str | int | Iterable[str | int] | None + The index to convert. + + Returns + ------- + List | int + List of integer indicies + + Raises + ------ + IndexError + _description_ + IndexError + _description_ + """ + # Convert to int, tuple(int), or None + if dims is None: out = None + elif isinstance(dims, str): out = self.get_index_from_name(dims) + elif isinstance(dims, int): + if dims > len(self): + raise IndexError(f'Axis index {dim} is out of bounds for signal with {self.ndim} dimensions.') + else: + out = dims if dims>=0 else len(self) + dims + elif isinstance(dims, Iterable): + out = [] + for dim in dims: + if isinstance(dim, str): + out.append(self.get_index_from_name(dim)) + elif isinstance(dim, int): + if dim > self.ndim: + raise IndexError(f'Axis index {dim} is out of bounds for signal with {self.ndim} dimensions.') + else: + out.append(dim if dim>=0 else len(self.dimensions) + dim) + return out + + def to_dict(self, + hidden: bool = False, + properties: bool = True, + exclude_keys: List[str] = [], + deep: bool = True + ): + """ + Recursively convert Dimension object to a dictionary. + + Parameters + ---------- + hidden : bool, optional + Include hidden attributes (starting with '_'), by default False. + properties : bool, optional + Include properties, by default True. + exclude_keys : List[str], optional + Keys to exclude from the dictionary, by default []. + deep : bool, optional + Recursively convert nested SEASerializable objects, by default True. + + Returns + ------- + dict + Dictionary representation of the Dimension object. + """ + return super().to_dict(hidden=hidden, properties=properties, exclude_keys=exclude_keys, deep=deep) + + def to_hdf5_group(self, parent_group: File|Group, + force_datasets: List = [], + name: str | None = None + ) -> None: + force_datasets = force_datasets#?+['dimensions'] + super().to_hdf5_group(parent_group=parent_group, + force_datasets=force_datasets, + name=name) + def to_sea(self, file_path: str, + force_datasets: List = []) -> None: + """Save dimension to HDF5.""" + force_datasets = force_datasets#?+['dimensions'] + super().to_sea(file_path, + force_datasets=force_datasets) + + def _get_tree_html(self, recursive_level: List[str] = 0, + exclude_keys: List[str] = [], + exclude_hidden: bool = True, + exclude_properties:bool = False, + promote_itterable_keys: List[str] = ['dimensions'] + ) -> str: + return super().get_tree_html(recursive_level, + exclude_keys=exclude_keys, + exclude_hidden=exclude_hidden, + exclude_properties=exclude_properties, + promote_itterable_keys=promote_itterable_keys + ) + def get_tree_html(self, recursive_level: int = 0, + exclude_keys: List[str] = [], + exclude_hidden: bool = True, + exclude_properties:bool = False, + promote_itterable_keys: List[str] = [] + ) -> str: + return self._get_tree_html(recursive_level, + exclude_keys=exclude_keys, + exclude_hidden=exclude_hidden, + exclude_properties=exclude_properties, + promote_itterable_keys=promote_itterable_keys + ['dimensions'] + ) + + def get_extents(self, kind: Literal['Axes','Image'] = 'Axes') -> List[float, List[float], List[Tuple[float]]]: + """Get the extents of all dimensions for plotting. + + Returns + ------- + List[float, Tuple[float]] + List of extents as [min, max] or list of (min, max) tuples for multi-D dimensions. + """ + extents = [dim.get_extent() for dim in self.dimensions] + hf_sc = [dim.scale/2 for dim in self.dimensions] + + if kind=='Image' and self.ndim==2: + # Flatten one level so extents becomes a single list like [xmin, xmax, ymax, ymin] + extents = [extents[1][0]-hf_sc[1], extents[1][1]+hf_sc[1], + extents[0][1]+hf_sc[0], extents[0][0]-hf_sc[0]] + return extents + +class Signal(SEASerializable): + + def __init__(self, data: NDArray|None = None, + name: str = 'Signal', + uuid: str = None, + dimensions: Dimensions|Dict|None = None, #BUG Not sure if Dict will work + signal_type: Literal['2D-EELS','1D-EELS','Diffraction','Image']|None = None, + dimensions_domain:Literal['local','global'] = 'local', + #? metadata_domain: Literal['local','global'] = 'local', + original_metadata: GeneralMetadata|None = None, + is_lazy: bool = False, + metadata: GeneralMetadata|None = None + ): + + self.data = data + self.name = name + self.uuid = uuid if uuid is not None else generate_uuid() + self.dimensions_domain = dimensions_domain + self.dimensions = dimensions + self._local_dimensions = dimensions + self.signal_type = signal_type + #self.metadata_domain = metadata_domain + self._original_metadata = original_metadata + self.is_lazy = is_lazy + self.metadata = metadata + + self._parent_SignalSet: SignalSet| None = None + + #HACK This is gross. should make GeneralMetadata sliceable like dict or with ints. Also not sure this will work when not available. + #? Might be worth thinking about where this should go. Could go in meta.instrument, even if it means all but instrument.detector is promoted, but this depends on how meta.instrument will be handled afer promotion. + self.detector: str = list(self.metadata.Instrument.Detectors.to_dict().keys())[0] if self.metadata is not None and hasattr(self.metadata, 'Instrument') and hasattr(self.metadata.Instrument, 'Detectors') else None + + @property + def dimensions(self): + if self.dimensions_domain=='global' and self._parent_SignalSet: + names_global = self._parent_SignalSet.dimensions.get_names() + names_local = self._local_dimensions.get_names() + axs_global_i = [names_global.index(nl) for nl in names_local if nl in names_global] + # axs_local_i = [names_local.index(ng) for ng in names_global if ng in names_local] + axs_local = [self._parent_SignalSet.dimensions.dimensions[i] for i in axs_global_i] + nav_local = [] + sig_local = [] + order_local = [] + for i in axs_global_i: + order_local.append(self._parent_SignalSet.dimensions.order.index(i)) + if i in self._parent_SignalSet.dimensions.nav_dimensions: + nav_local.append(names_local.index(names_global[i])) + if i in self._parent_SignalSet.dimensions.sig_dimensions: + sig_local.append(names_local.index(names_global[i])) + if len(nav_local)==0: nav_local = None + if len(sig_local)==0: sig_local = None + self._local_dimensions = Dimensions(axs_local, nav_dimensions=nav_local, sig_dimensions=sig_local) + return self._local_dimensions + @dimensions.setter + def dimensions(self, dimensions): + self._local_dimensions = dimensions + if self.dimensions_domain=='global' and self._parent_SignalSet: + warn('This signal is using global dimensions. Setting the the signal locally will set the calibration_domian to local.') + while True: + user_input = input("Do you want to continue? (Y/N): ").upper() # Convert to uppercase for case-insensitivity + if user_input == 'Y': + print("Proceeding...") + # Add the code to execute if the user chooses 'Y' here + break # Exit the loop once a valid 'Y' is entered + elif user_input == 'N': + print("Exiting...") + # Add the code to execute if the user chooses 'N' here + break # Exit the loop once a valid 'N' is entered + else: + print("Invalid input. Please enter 'Y' or 'N'.") + self.dimensions_domain = 'local' + self._local_dimensions = dimensions + else: + self._local_dimensions = dimensions + + @property + def original_metadata(self) -> GeneralMetadata: + return self._original_metadata + @original_metadata.setter + def original_metadata(self, values: Any) -> UserWarning | None: + if self._original_metadata is None: self._original_metadata = values + else: raise UserWarning('original_metadata is read-only. If it is necessary to change the original metadata use `_original_metadata` to set the value, but we do recomend against changing such values.') + + def __str__(self): + return f'{self.name}-signal' + + def __repr__(self): + return f'' + + def __getitem__(self, key: Union[int, float, slice, tuple ]): + """Support indexing and slicing of the Signal. + + Parameters + ---------- + key : Union[int, float, slice, tuple, EllipsisType] + Index specification. Can include integers, floats (converted to nearest index), + slices, ellipsis, or tuples of these. + - int: direct indexing + - float: converted to nearest index using dimension calibration + - slice: regular Python slicing, can include float values + - Ellipsis: expands to cover remaining dimensions + + ToDo + ---- + """ + if not isinstance(key, tuple): key = (key,) + + # Track which dimensions remain and their new sizes + remaining_dims = self.dimensions.dimensions.copy() + new_sizes = [] + + # Handle Ellipsis expansion + n_indices = len(key) + n_dims = len(self.dimensions.dimensions) + ellipsis_pos = None + + #for i, k in enumerate(key): + # if k is Ellipsis: + # ellipsis_pos = i + # break + + if ellipsis_pos is not None: + # Calculate how many dimensions the Ellipsis represents + n_extra = n_dims - (n_indices - 1) + # Replace Ellipsis with appropriate number of slice(None) + expanded_key = key[:ellipsis_pos] + (slice(None),) * n_extra + key[ellipsis_pos + 1:] + else: + expanded_key = key + + # Ensure we don't have too many indices + if len(expanded_key) > n_dims: + raise IndexError(f'Too many indices: array is {n_dims}-dimensional, but {len(expanded_key)} were indexed') + + # Pad with full slices if we have too few indices + if len(expanded_key) < n_dims: + expanded_key = expanded_key + (slice(None),) * (n_dims - len(expanded_key)) + + key_reg = tuple() + # Now process each key with the proper dimension + for i, k in enumerate(expanded_key): + orig_dim_i = self.dimensions.order[i] + orig_dim = self.dimensions.dimensions[orig_dim_i] + if isinstance(k, int): + remaining_dims.remove(orig_dim) + elif isinstance(k, float): + k = orig_dim.find_nearest_index(k) + # This dimension is removed by indexing + remaining_dims.remove(orig_dim) + elif isinstance(k, slice): + start = k.start + stop = k.stop + step = k.step + if isinstance(start, float): start = orig_dim.find_nearest_index(start) + if isinstance(stop, float): stop = orig_dim.find_nearest_index(stop) + if isinstance(step, float): step = int(round(step / orig_dim.scale, 0)) + k = slice(start, stop, step) + # This dimension remains but might have a new size + # Calculate new size for sliced dimension + slice_indices = k.indices(orig_dim.size) + new_size = len(range(*slice_indices)) + new_sizes.append(new_size) + else: + # Ellipse do not remove + pass + key_reg += (k,) + + sliced_data = self.data[key_reg] + + # Create new signal with updated dimensions + new_signal = self.deepcopy() + new_signal.data = sliced_data + + # Update dimensions for the new signal + if remaining_dims: + new_dims = Dimensions( + dimensions=[dim.deepcopy() for dim in remaining_dims], + nav_dimensions=[i for i, dim in enumerate(remaining_dims) + if dim in [self.dimensions[d] for d in self.dimensions.nav_dimensions]], + sig_dimensions=[i for i, dim in enumerate(remaining_dims) + if dim in [self.dimensions[d] for d in self.dimensions.sig_dimensions]], + ) + # Update sizes for sliced dimensions + for dim, new_size in zip(new_dims.dimensions, new_sizes): + dim.size = new_size + new_signal.dimensions = new_dims + + return new_signal + + #TODO: Implement a wrapper for numpy functions that modifies the dimensions accordingly + def __array__(self, dtype:DTypeLike=None) -> Any: + """Allow numpy to treat this object as an array.""" + if dtype is None: + return self.data + return self.data.astype(dtype) + + def __array_function__(self, func, types, args, kwargs): + """Handle numpy array functions like sum(), mean(), etc. + + This is called by numpy when array-like objects are passed to numpy functions. + It differs from __array_ufunc__ which handles element-wise operations. + """ + # Get the actual array data from any Signal objects + arrays = [] + signal_inputs = [] + for arg in args: + if isinstance(arg, Signal): + arrays.append(arg.data) + signal_inputs.append(arg) + else: + arrays.append(arg) + + # Convert any dimension references and get the dimension key being used + dim_key = check_dimensions_call(kwargs, func) + if dim_key: + kwargs[dim_key] = self.dimensions.get_dims_as_int(kwargs[dim_key]) + if isinstance(kwargs[dim_key], List): kwargs[dim_key] = tuple(kwargs[dim_key]) + + if isinstance(kwargs[dim_key], int | str | Iterable): dims_to_remove = np.atleast_1d(kwargs[dim_key]) + else: raise KeyError(f'{dim_key} takes int, str, or an Iterable and {kwargs[dim_key]} was provided.') + remaining_dims = {i: dim for i, dim in enumerate(self.dimensions.dimensions) + if i not in dims_to_remove} + + # Call the numpy function with our data arrays + result = func(*arrays, **kwargs) + + # If the result is an array, wrap it in a Signal + if isinstance(result, np.ndarray): + if np.ndim(result) == self.data.ndim: + result = self.deepcopy_with_new_data(result) + elif np.ndim(result) < self.data.ndim: + result = self.deepcopy_with_reduced_data_dim(data=result, keep_dim=remaining_dims) + return result + + def __array_ufunc__(self, ufunc: np.ufunc, method: str, *inputs, **kwargs) -> Any: + """Handle numpy universal functions (ufuncs). + + Parameters + ---------- + ufunc : np.ufunc + A numpy function that is applied to the input data. + method : str + How the ufunc operates on the inputs. This partially handles dimensionality expansion and contraction. + *inputs : tuple + Input arrays to the ufunc. + **kwargs : dict + Additional keyword arguments to pass to ufunc. + """ + # Convert any Signal objects in inputs to their data arrays + args = [] + signal_inputs = [] + for i in inputs: + if isinstance(i, Signal): + args.append(i.data) + signal_inputs.append(i) + else: + args.append(i) + + # Check that a signal exists in the inputs + if not signal_inputs: raise TypeError('At least one input must be a Signal object. I am not sure how this was called internally without a Signal. Please submit a detailed bug report.') + + # Catch axis indexing by name + possible_dim_keys = ['axis', 'axes', 'dim'] + dim_key = [k for k in possible_dim_keys if k in kwargs.keys()] + if len(dim_key) > 1: + raise ValueError(f'Only one of {possible_dim_keys} can be used at a time, but {dim_key} were provided. Check what dimensional refference {ufunc} permits and only provide that refference.') + elif len(dim_key) == 1: + dim_key = dim_key[0] + else: + dim_key = None + + if dim_key: + # dim is used in SEA and PyTorch but not in numpy. + # Catch when SEA uses dim and convert to the funcitons allowed dim representation. + if dim_key=='dim': + for alwd_dim_key in possible_dim_keys[:-1]: + if alwd_dim_key in signature(ufunc).paramaters: + kwargs[alwd_dim_key] = kwargs['dim'] + dim_key = alwd_dim_key + del kwargs['dim'] + break + + # Convert to int, tuple(int), or None + if kwargs[dim_key] is None: dims = None + elif isinstance(kwargs[dim_key], str): + dims = signal_inputs[0].dimensions.get_index_from_name(kwargs[dim_key]) + elif isinstance(kwargs[dim_key], int): + if kwargs[dim_key] > signal_inputs[0].dimensions.ndim: + raise IndexError(f'Axis index {dim} is out of bounds for signal with {signal_inputs[0].dimensions.ndim} dimensions.') + else: + dim = kwargs[dim_key] + dims = dim if dim>=0 else len(self.dimensions) + dim + elif isinstance(kwargs[dim_key], Iterable): + dims = tuple() + for dim in kwargs[dim_key]: + if isinstance(dim, str): + dims += (signal_inputs[0].dimensions.get_index_from_name(dim), ) + elif isinstance(dim, int): + if dim > signal_inputs[0].dimensions.ndim: + raise IndexError(f'Axis index {dim} is out of bounds for signal with {signal_inputs[0].dimensions.ndim} dimensions.') + else: + dim = dim if dim>=0 else len(self.dimensions) + dim + dims += (dim, ) + kwargs[dim_key] = dims + + # Call the ufunc on the underlying data + result = getattr(ufunc, method)(*args, **kwargs) + + # If result is an array, wrap it back in a Signal + if isinstance(result, np.ndarray): + # Use the first Signal's attributes for the result + if signal_inputs: + base_signal = signal_inputs[0] + out = base_signal.deepcopy() + out.data = result + + # Update dimensions based on the ufunc operation + if result.ndim != base_signal.data.ndim: + # Dimension reduction occurred + remaining_dims = {i: n for i, n in enumerate(base_signal.dimensions)} + + if dim_key: + if isinstance(kwargs[dim_key], int): kwargs[dim_key] = tuple([kwargs[dim_key]]) + for dim in kwargs[dim_key]: remaining_dims.pop(dim) + + new_dims = Dimensions( + dimensions=remaining_dims.values(), + nav_dimensions=[i for i, dim in enumerate(remaining_dims) + if dim in [base_signal.dimensions[d] for d in base_signal.dimensions.nav_dimensions]], + sig_dimensions=[i for i, dim in enumerate(remaining_dims) + if dim in [base_signal.dimensions[d] for d in base_signal.dimensions.sig_dimensions]], + ) + out.dimensions = new_dims + + return out + + # Return scalars or other types as-is + return result + + def deepcopy_with_new_data(self, data:np.ndarray): + """Copy the signal with new data.""" + out = self.deepcopy() + out.data = data + return out + def deepcopy_with_reduced_data_dim(self, data:np.ndarray, keep_dim:Tuple[int|str]) -> Signal: #! + """Copy the signal with new data of reduced dimensionality.""" + out = self.deepcopy() + out.data = data + + # Update dimensions based on the ufunc operation + new_dims = Dimensions( + dimensions=[out.dimensions[i] for i in keep_dim], + nav_dimensions=[i for i, dim in enumerate(keep_dim) + if dim in [self.dimensions[d] for d in self.dimensions.nav_dimensions]], + sig_dimensions=[i for i, dim in enumerate(keep_dim) + if dim in [self.dimensions[d] for d in self.dimensions.sig_dimensions]], + ) + out.dimensions = new_dims + + return out + + def to_dict(self, + hidden: bool = False, + properties: bool = True, + exclude_keys: List[str] = ['data'], + deep: bool = True + ): + """ + Recursively convert Dimension object to a dictionary. + + Parameters + ---------- + hidden : bool, optional + Include hidden attributes (starting with '_'), by default False. + properties : bool, optional + Include properties, by default True. + exclude_keys : List[str], optional + Keys to exclude from the dictionary, by default []. + deep : bool, optional + Recursively convert nested SEASerializable objects, by default True. + + Returns + ------- + dict + Dictionary representation of the Dimension object. + """ + return super().to_dict(hidden=hidden, properties=properties, exclude_keys=exclude_keys, deep=deep) + + def to_hdf5_group(self, parent_group: File|Group, + force_datasets: List = ['data'], + name: str | None = None + ) -> None: + super().to_hdf5_group(parent_group=parent_group, + force_datasets=force_datasets, + name=name) + def to_sea(self, file_path: str, + force_datasets: List = ['data']) -> None: + """Save dimension to HDF5.""" + super().to_sea(file_path=file_path, + force_datasets=force_datasets) + + def _get_tree_html(self, recursive_level: List[str] = 0, + exclude_keys: List[str]= ['data', 'original_metadata'], + exclude_hidden: bool = True, + exclude_properties:bool = False, + promote_itterable_keys: List[str] = [] + ) -> str: + return super().get_tree_html(recursive_level, + exclude_keys=exclude_keys, + exclude_hidden=exclude_hidden, + exclude_properties=exclude_properties, + promote_itterable_keys=promote_itterable_keys + ) + def get_tree_html(self, recursive_level: int = 0, + exclude_keys: List[str] = [], + exclude_hidden: bool = True, + exclude_properties:bool = False, + promote_itterable_keys: List[str] = [] + ) -> str: + return self._get_tree_html(recursive_level, + exclude_keys=exclude_keys+['data', 'original_metadata'], + exclude_hidden=exclude_hidden, + exclude_properties=exclude_properties, + promote_itterable_keys=promote_itterable_keys + ) + + def infer_plot_dims(self,dims,fnc): + if dims is None: + if self.dimensions.ndim == 2: dims = [0,1] # The signal is already plottable + elif self.dimensions.ndim == 1: dims = [0] # The signal is already plottable + elif 0 < len(self.dimensions.sig_dimensions) <= 2: dims = self.dimensions.sig_dimensions # Plot the signal + elif 0 < len(self.dimensions.nav_dimensions) <= 2: dims = self.dimensions.nav_dimensions # Plot the navigation + else: raise ValueError('The total, signal, and navigation dimensions are larger than 2 dimensions so a dimension could not be infered. Ploting >2D is not yet implemented') + elif dims == 'sig': + if len(self.dimensions.sig_dimensions) <=2: dims = self.dimensions.sig_dimensions + else: raise ValueError('The signal dimensions being plotted is larger than 2 dimensions, which is not yet implemented') + elif dims == 'nav': + if len(self.dimensions.nav_dimensions) <=2: dims = self.dimensions.nav_dimensions + else: raise ValueError('The navigation dimensions being plotted is larger than 2 dimensions, which is not yet implemented') + + #TODO use self.dimensions.get_dims_as_int() + if len(dims)>2: raise ValueError('The signal dimensions being plotted is larger than 2 dimensions, which is not yet implemented') + else: dims = self.dimensions.get_dims_as_int(dims) + + dims_remain = tuple(i for i in range(len(self.dimensions)) if i not in dims) + + if len(dims_remain) == 0: + return self,dims + else: + return fnc(self, axis=dims_remain),dims + + def show(self, + ax: mplAxes|None = None, + dims: None|Literal['sig','nav']|Iterable[int, str]= None, + fnc: Callable|None = np.sum, + filename: str|None = None, + **kwargs + ) -> Any: + + if not PLOTTING_AVAILABLE: + raise ImportError("Plotting functionality requires sea-eco plotting module. Use matplotlib directly or install sea-eco.") + + sig,dims = self.infer_plot_dims(dims,fnc) + + if 'xlabel' not in kwargs: + kwargs['xlabel'] = f'{sig.dimensions[-1].name} ({sig.dimensions[-1].units})' + if len(dims)==2: + if 'extent' not in kwargs: kwargs['extent'] = sig.dimensions.get_extents(kind='Image') + if 'ylabel' not in kwargs: kwargs['ylabel'] = f'{sig.dimensions[-2].name} ({sig.dimensions[-2].units})' + if 'scale_bar_kwargs' not in kwargs: kwargs['scale_bar_kwargs'] = {'units':sig.dimensions[-1].units} + else: + if 'units' not in kwargs['scale_bar_kwargs']: + kwargs['scale_bar_kwargs']['units'] =sig.dimensions[-1].units + if len(dims)==1: + if 'ylabel' not in kwargs: kwargs['xlabel'] = f'{sig.dimensions[0].name} ({sig.dimensions[0].units})' + p = plot_nd_array(sig.data, ax=ax, **kwargs) + if filename is not None: + save_plot(filename) + return p + + def image(self, + dims: None|Literal['sig','nav']|Iterable[int, str]= None, + fnc: Callable|None = np.sum, + filename: str|None = None ) -> None: + + if not PLOTTING_AVAILABLE: + raise ImportError("Plotting functionality requires sea-eco plotting module. Use matplotlib directly or install sea-eco.") + + sig,dims = self.infer_plot_dims(dims,fnc) + size = np.asarray( [ d.get_extent() for d in sig.dimensions ] ) + size = size[:,1]-size[:,0] + units = [ d.units for d in sig.dimensions ] + save_image(sig.data,size,units[0],filename=filename) + +class SignalSet(SEASerializable): + + def __init__(self, + signals: Iterable[Signal]|None = None, main_signal: int|None = None, + name: str|None = None, + metadata: GeneralMetadata|None = None, + dimensions: Dimensions|None = None, + uuid: str|None = None, + merge_dimensions:bool = False, merge_metadata: bool = False + ) -> None: + + #Initialize kwargs + self.uuid = uuid if uuid is not None else generate_uuid() + self.main_signal = main_signal + self.signals = [] + self.dimensions = dimensions if dimensions is not None else Dimensions() + self.metadata = metadata + + # Add signals + if signals is not None: # loop signals and add them + if main_signal is None: main_signal=0 #if main_signal is not defined define the first element as the main signal + else: signals.insert(0, signals.pop(main_signal)) # move the main signal to the first position in the list + self.main_signal = 0 # redefine main_signal as the first signal, as we just inforce + + for signal in signals: + signal._parent_SignalSet = self # First set parent reference so dimensions can be accessed + self.add_signal(signal, merge_metadata=merge_metadata, merge_dimensions=merge_dimensions) # Then add signals with metadata and dimensions merging + + def __getitem__(self, key:int|str): + if isinstance(key, int): return self.signals[key] + elif isinstance(key, str): return self.signals[self.get_index_from_name(key)] + else: raise TypeError(f'Only integers and strings are allowed but a {type(key)} was provided.') + + # def to_hdf5_grop( + # def to_sea(self, group, name): + # pass + + def _get_tree_html(self, recursive_level: List[str] = 0, + exclude_keys: List[str] = [], + exclude_hidden: bool = True, + exclude_properties:bool = False, + promote_itterable_keys: List[str] = ['signals'] + ) -> str: + return super().get_tree_html(recursive_level, + exclude_keys=exclude_keys, + exclude_hidden=exclude_hidden, + exclude_properties=exclude_properties, + promote_itterable_keys=promote_itterable_keys + ) + def get_tree_html(self, recursive_level: int = 0, + exclude_keys: List[str] = [], + exclude_hidden: bool = True, + exclude_properties:bool = False, + promote_itterable_keys: List[str] = [] + ) -> str: + return self._get_tree_html(recursive_level, + exclude_keys=exclude_keys, + exclude_hidden=exclude_hidden, + exclude_properties=exclude_properties, + promote_itterable_keys=promote_itterable_keys + ['signals'] + ) + + def add_signal(self, signal, + merge_metadata:bool=False, + merge_dimensions:bool=False, + meta_kwargs=dict(kind='append', warn_duplicate=False, inform_new=False)): + signal = signal.deepcopy() #? Should this be coppied? + if merge_metadata: + if self.metadata is None: self.metadata = signal.metadata + elif signal.metadata is not None: self.metadata.merge(signal.metadata, **meta_kwargs) + + if merge_dimensions: + # dimensions is initialized in __init__, so this case shouldn't happen anymore + if len(self.dimensions.nav_dimensions)==0 and len(signal.dimensions.nav_dimensions)!=0: + self.dimensions.nav_dimensions = signal.dimensions.nav_dimensions + if len(self.dimensions.sig_dimensions)==0 and len(signal.dimensions.sig_dimensions)!=0: + self.dimensions.sig_dimensions = signal.dimensions.sig_dimensions + if len(self.dimensions.nav_dimensions)==0 and len(signal.dimensions.nav_dimensions)!=0: + self.dimensions.nav_dimensions = signal.dimensions.nav_dimensions + for ax in signal.dimensions.dimensions: + if ax.name not in self.dimensions.get_names(): + self.dimensions.add_dimension(ax) + signal.dimensions_domain = 'global' + self.signals.append(signal) + + def get_names(self) -> List[str]: + return [sig.name for sig in self.signals] + + def get_index_from_name(self, name:str) -> int: + names = self.get_names() + if name in names: return names.index(name) + else: raise KeyError(f'A key of {name} was provided and the signal names are {names}') + +class AcquisitionSet(SignalSet): + + def __init__(self, + signals: Iterable[Signal]|None = None, main_signal: int|None = None, + metadata: GeneralMetadata|None = None, + dimensions: Dimensions|None = None, + merge_dimensions: bool=True, merge_metadata: bool = True, + instrument_uuid: str|None = None + ) -> None: + + super().__init__(signals=signals, main_signal=main_signal, + metadata=metadata, + dimensions=dimensions, + merge_dimensions=merge_dimensions, merge_metadata=merge_metadata) + if self._check_aquisition_uuids(): + self.uuid = self.signals[0].metadata.Instrument.Scan.scan_uuid + else: + self.uuid = generate_uuid() + self.instrument_uuid = instrument_uuid + + def _check_aquisition_uuids(self): + uuids = [sig.metadata.Instrument.Scan.scan_uuid for sig in self.signals] + if len(self.signals)==0: + return False + elif len(set(uuids))!=1: + warn('Not all signals in this set have the same scan uuid.') + return False + return True diff --git a/src/pyslice/data/__init__.py b/src/pyslice/data/NO__init__.py similarity index 100% rename from src/pyslice/data/__init__.py rename to src/pyslice/data/NO__init__.py diff --git a/src/pyslice/data/pyslice_serial.py b/src/pyslice/data/pyslice_serial.py new file mode 100644 index 0000000..097880d --- /dev/null +++ b/src/pyslice/data/pyslice_serial.py @@ -0,0 +1,243 @@ +""" +PySliceSerial mixin for HDF5/SEA serialization of PySlice data classes. + +This mixin provides generalized serialization/deserialization for classes +that inherit from Signal but have special attributes (tensors, Paths, etc.) +that need conversion for HDF5 storage. +""" +import numpy as np +from pathlib import Path +from h5py import File, Group, Dataset + +try: + import sys + sys.path.insert(1,"../../") + from pySEA.sea_eco.architecture.base_structure_numpy import Signal, Dimensions, Dimension, Metadata, safe_decode +except Exception as e: + class Signal: + def to_sea(self,*args,**kwargs): + print("ERROR: pySEA import failed, this functionality is not available") + Dimensions,Dimension,Metadata = None,None,None + print("failed to import pySEA:",e) + +def _to_numpy(x): + """Convert tensor or array-like to numpy array.""" + if x is None: + return None + if hasattr(x, 'cpu'): + return x.cpu().numpy() + return np.asarray(x) + + +class PySliceSerial: + """ + Mixin class providing generalized HDF5/SEA serialization for PySlice data classes. + + Subclasses should define a `_sea_config` dict with the following optional keys: + + - tensor_attrs: List of attribute names that are torch tensors (converted to numpy) + - path_attrs: List of attribute names that are Path objects (converted to string) + - tuple_list_attrs: List of attribute names that are lists of tuples (converted to arrays) + - exclude_attrs: List of attribute names to exclude from serialization + - force_datasets: List of attribute names to store as HDF5 datasets (not attrs) + - default_attrs: Dict of default values to set during deserialization + + Example: + class MyData(PySliceSerial, Signal): + _sea_config = { + 'tensor_attrs': ['_array', '_kxs', '_kys'], + 'path_attrs': ['cache_dir'], + 'tuple_list_attrs': ['probe_positions'], + 'exclude_attrs': ['probe', '_wf_array'], + 'force_datasets': ['_array', 'probe_positions'], + } + """ + + _sea_config = {} + + def to_hdf5_group(self, parent_group, force_datasets=None, name=None): + """Serialize to HDF5 group with automatic type conversions.""" + config = getattr(self, '_sea_config', {}) + + tensor_attrs = config.get('tensor_attrs', []) + path_attrs = config.get('path_attrs', []) + tuple_list_attrs = config.get('tuple_list_attrs', []) + exclude_attrs = config.get('exclude_attrs', []) + config_force_datasets = config.get('force_datasets', []) + + if force_datasets is None: + force_datasets = ['data'] + force_datasets = list(force_datasets) + config_force_datasets + + # Store originals for restoration + originals = {} + + # Convert tensor attributes to numpy + for attr in tensor_attrs: + if hasattr(self, attr): + originals[attr] = getattr(self, attr) + setattr(self, attr, _to_numpy(getattr(self, attr))) + # Ensure tensor attrs are stored as datasets + if attr not in force_datasets: + force_datasets.append(attr) + + # Convert Path attributes to string + for attr in path_attrs: + if hasattr(self, attr): + originals[attr] = getattr(self, attr) + val = getattr(self, attr) + setattr(self, attr, str(val) if val is not None else None) + + # Convert list of tuples to numpy array + for attr in tuple_list_attrs: + if hasattr(self, attr): + originals[attr] = getattr(self, attr) + val = getattr(self, attr) + if val is not None: + setattr(self, attr, np.array(val)) + if attr not in force_datasets: + force_datasets.append(attr) + + # Temporarily remove non-serializable attributes + for attr in exclude_attrs: + if hasattr(self, attr): + originals[attr] = getattr(self, attr) + try: + delattr(self, attr) + except AttributeError: + setattr(self, attr, None) + + try: + # Call parent's to_hdf5_group (Signal's method) + result = super().to_hdf5_group(parent_group, force_datasets=force_datasets, name=name) + finally: + # Restore original attributes + for attr, val in originals.items(): + setattr(self, attr, val) + + return result + + def to_sea(self, file_path, force_datasets=None): + """Save to .sea file with automatic type conversions.""" + config = getattr(self, '_sea_config', {}) + config_force_datasets = config.get('force_datasets', []) + + if force_datasets is None: + force_datasets = ['data'] + force_datasets = list(force_datasets) + config_force_datasets + + super().to_sea(file_path, force_datasets=force_datasets) + + def from_hdf5_group(self, group): + """Deserialize from HDF5 group with automatic type conversions.""" + #from .signal import safe_decode, Dimensions, GeneralMetadata + + config = getattr(self, '_sea_config', {}) + + tensor_attrs = config.get('tensor_attrs', []) + path_attrs = config.get('path_attrs', []) + tuple_list_attrs = config.get('tuple_list_attrs', []) + exclude_attrs = config.get('exclude_attrs', []) + default_attrs = config.get('default_attrs', {}) + + # Build mapping from storage keys (without _) to internal attr names + key_map = {} + for attr in tensor_attrs + path_attrs + tuple_list_attrs: + storage_key = attr[1:] if attr.startswith('_') else attr + key_map[storage_key] = attr + + # Initialize default attributes (Signal expects these) + signal_defaults = { + '_original_metadata': None, + '_parent_SignalSet': None, + 'detector': None, + 'is_lazy': False, + '_array': None, + } + for attr, val in signal_defaults.items(): + setattr(self, attr, val) + + # Initialize excluded attrs to None + for attr in exclude_attrs: + setattr(self, attr, None) + + # Initialize user-specified defaults + for attr, val in default_attrs.items(): + setattr(self, attr, val) + + # Read datasets (large arrays) + for key, item in group.items(): + if isinstance(item, Dataset): + value = item[()] + attr_name = key_map.get(key, key) + # Check for private version + if not hasattr(self, attr_name) and hasattr(self, f'_{key}'): + attr_name = f'_{key}' + setattr(self, attr_name, value) + + # Read attributes (scalars/small values) + for key, val in group.attrs.items(): + if key == 'sea_type': + continue + decoded_val = safe_decode(val) + attr_name = key_map.get(key, key) + + # Check for private version + if not hasattr(self, attr_name): + if hasattr(self, f'_{key}'): + attr_name = f'_{key}' + + setattr(self, attr_name, decoded_val) + + # Handle nested groups (Dimensions, GeneralMetadata) + for key, item in group.items(): + if isinstance(item, Group): + sea_type = safe_decode(item.attrs.get('sea_type', b'')) + if sea_type == 'Dimensions': + dims = Dimensions() + dims.from_hdf5_group(item) + self._local_dimensions = dims + elif sea_type == 'GeneralMetadata': + meta = GeneralMetadata() + meta.from_hdf5_group(item) + if key == 'metadata' or key == 'Metadata': + self.metadata = meta + + # Post-process: convert types back + for attr in path_attrs: + if hasattr(self, attr): + val = getattr(self, attr) + if val is not None and isinstance(val, str): + setattr(self, attr, Path(val)) + + for attr in tuple_list_attrs: + if hasattr(self, attr): + val = getattr(self, attr) + if val is not None and isinstance(val, np.ndarray): + setattr(self, attr, [tuple(p) for p in val]) + + @classmethod + def load(cls, file_path): + """ + Load an object from a .sea file. + + Args: + file_path: Path to the .sea file + + Returns: + Instance of the class populated with data from the file + """ + file_path = Path(file_path) + if file_path.suffix != '.sea': + file_path = file_path.with_suffix('.sea') + + # Create empty instance without calling __init__ + obj = cls.__new__(cls) + + with File(file_path, 'r') as f: + if len(f) != 1: + raise ValueError("The HDF5 file contains multiple groups.") + main_group = f[list(f.keys())[0]] + obj.from_hdf5_group(main_group) + + return obj diff --git a/src/pyslice/md/molecular_dynamics.py b/src/pyslice/md/molecular_dynamics.py index e464f95..81c4bed 100644 --- a/src/pyslice/md/molecular_dynamics.py +++ b/src/pyslice/md/molecular_dynamics.py @@ -128,7 +128,7 @@ class MDCalculator: Integrates with PySlice's Trajectory and Loader infrastructure. """ - def __init__(self, model_name: str = 'orb-v3-direct-inf-omat', device: str = 'cpu', precision: str = 'float32-high'): + def __init__(self, model_name: str = 'orb-v3-direct-inf-omat', device: str = 'cpu', precision: str = 'float32-high', weights_path: Optional[str] = None): """ Initialize MD calculator. @@ -136,10 +136,12 @@ def __init__(self, model_name: str = 'orb-v3-direct-inf-omat', device: str = 'cp model_name: ORB model to use device: Device for ORB calculations ('cpu', 'cuda', etc.) precision: Precision for ORB calculations ('float32-high', 'float32-highest', etc.) + weights_path: Optional local path to model weights (bypasses network download) """ self.model_name = model_name self.device = device self.precision = precision + self.weights_path = weights_path self.calculator = None # Available ORB models @@ -169,15 +171,22 @@ def _setup_calculator(self) -> bool: model_func = getattr(pretrained, model_func_name) + # Build kwargs for model loading + model_kwargs = {'precision': self.precision, 'compile': False} + if self.weights_path: + logger.info(f"Using local weights: {self.weights_path}") + model_kwargs['weights_path'] = self.weights_path + # ORB doesn't natively support MPS, so we load on CPU then move if self.device == 'mps': logger.info("MPS detected: loading on CPU, converting to float32, moving to MPS...") # compile=False required for MPS (dynamo has float64 issues) - orbff = model_func(device='cpu', precision=self.precision, compile=False) + orbff = model_func(device='cpu', **model_kwargs) orbff = orbff.float() # Ensure float32 (MPS doesn't support float64) orbff = orbff.to('mps') else: - orbff = model_func(device=self.device, precision=self.precision) + # compile=False to avoid slow torch.compile on first run + orbff = model_func(device=self.device, **model_kwargs) self.calculator = ORBCalculator(orbff) @@ -552,7 +561,8 @@ def run_production( # Convert to PySlice Trajectory automatically # Use ASE directly to preserve velocities (OVITO strips them) - timestep_ps = self.timestep / 1000.0 # Convert fs to ps + # Account for save_interval: actual time between frames = timestep * save_interval + timestep_ps = self.timestep * self.save_interval / 1000.0 # Convert fs to ps logger.info(f"Converting trajectory to PySlice format...") from ase.io import read as aseread diff --git a/src/pyslice/multislice/calculators.py b/src/pyslice/multislice/calculators.py index 3bd5b24..2fc4537 100644 --- a/src/pyslice/multislice/calculators.py +++ b/src/pyslice/multislice/calculators.py @@ -3,7 +3,7 @@ import logging from typing import Optional, Tuple, List from tqdm import tqdm -import time +import time,os import hashlib try: @@ -34,7 +34,7 @@ from .trajectory import Trajectory from ..postprocessing.wf_data import WFData from .sed import SED -from ..backend import zeros,expand_dims +from ..backend import zeros,expand_dims,to_cpu logger = logging.getLogger(__name__) @@ -165,16 +165,32 @@ def setup( self.dx = xs[1]-xs[0] ; self.dy = ys[1]-ys[0] ; self.dy = ys[1]-ys[0] # calculate kxs kys here, so we can crop them, since we'll pre-allocate wavefunction_data below - self.kxs = xp.fft.fftshift(xp.fft.fftfreq(self.nx, self.sampling)) # k-space in 1/Å - self.kys = xp.fft.fftshift(xp.fft.fftfreq(self.ny, self.sampling)) # k-space in 1/Å - self.i1 = xp.argwhere(self.kxs >= -max_kx)[0][0] # first element >= - self.i2 = xp.argwhere(self.kxs <= max_kx)[-1][0]+1 # last element <=, +1, so i1:i2 includes i2 - self.j1 = xp.argwhere(self.kys >= -max_ky)[0][0] - self.j2 = xp.argwhere(self.kys <= max_ky)[-1][0]+1 - self.kxs = self.kxs[self.i1:self.i2] - self.kys = self.kys[self.j1:self.j2] + self.kxs_uncrop = xp.fft.fftshift(xp.fft.fftfreq(self.nx, self.sampling)) # k-space in 1/Å + self.kys_uncrop = xp.fft.fftshift(xp.fft.fftfreq(self.ny, self.sampling)) # k-space in 1/Å + self.i1 = xp.argwhere(self.kxs_uncrop >= -max_kx)[0][0] # first element >= + self.i2 = xp.argwhere(self.kxs_uncrop <= max_kx)[-1][0]+1 # last element <=, +1, so i1:i2 includes i2 + self.j1 = xp.argwhere(self.kys_uncrop >= -max_ky)[0][0] + self.j2 = xp.argwhere(self.kys_uncrop <= max_ky)[-1][0]+1 + self.kxs = self.kxs_uncrop[self.i1:self.i2] + self.kys = self.kys_uncrop[self.j1:self.j2] self.nx = self.i2 - self.i1 ; self.ny = self.j2 - self.j1 ; nx = self.nx ; ny = self.ny + # Preferred to pass probe_xs and probe_ys from which we will define a grid + if self.probe_xs is not None and self.probe_ys is not None: + x,y = np.meshgrid(self.probe_xs,self.probe_ys,indexing='ij') + self.probe_positions = np.asarray(list(zip(x.flat,y.flat))) + + # If probe_positions provided but not probe_xs/probe_ys, derive them + elif self.probe_positions is not None: + positions = np.asarray(self.probe_positions) + self.probe_xs = sorted(list(set(positions[:, 0]))) + self.probe_ys = sorted(list(set(positions[:, 1]))) + + # Set up default probe position if not provided + if self.probe_positions is None: + self.probe_positions = [(lx/2, ly/2)] # Center probe + self.probe_xs = [lx/2] ; self.probe_ys = [ly/2] + # Create probe on the correct device from the start self.base_probe = Probe(xs, ys, self.aperture, self.voltage_eV, device=self.device, probe_xs=self.probe_xs, probe_ys=self.probe_ys, probe_positions=self.probe_positions) self.base_probe.applyShifts() @@ -195,7 +211,30 @@ def setup( self.complex_dtype = np.complex128 self.float_dtype = np.float64 - + def preview_probes(self): + positions = self.trajectory.positions[0] + atom_types = self.trajectory.atom_types + atom_type_names = [] + for atom_type in atom_types: + if atom_type in self.element_map: + atom_type_names.append(self.element_map[atom_type]) + else: + atom_type_names.append(atom_type) + potential = Potential(self.xs, self.ys, self.zs, positions, atom_type_names, kind="kirkland", device=self.device, slice_axis=self.slice_axis) + potential.build() + potential.flatten() + import matplotlib.pyplot as plt + fig, ax = plt.subplots() + array = np.absolute(to_cpu(potential.array))[:,::-1,0].T # imshow convention: y,x. our convention: x,y, and flip y (0,0 upper-left) + xs = to_cpu(potential.xs) ; ys = to_cpu(potential.ys) + extent = (np.amin(xs),np.amax(xs),np.amin(ys),np.amax(ys)) + print(extent) + ax.imshow(array, cmap="inferno", extent=extent) + ax.set_xlabel("x ($\\AA$)") ; ax.set_ylabel("y ($\\AA$)") + pp = np.asarray(self.base_probe.probe_positions) + ax.scatter(pp[:,0],pp[:,1],c='r') + plt.show() + def run(self) -> WFData: @@ -231,6 +270,12 @@ def run(self) -> WFData: # Show detailed progress for single-frame runs show_progress = (frame_idx == 0 and self.n_frames == 1) + # special case: no frames cached, but we clearly finished and got to tacaw. if so, don't bother regenerating + # this allows cache_levels = [] to be used for disk space savings + if os.path.exists( self.output_dir / f"tacaw.npy" ) and not os.path.exists( cache_file ): + pbar.update(1) + continue + positions = self.trajectory.positions[frame_idx] atom_types = self.trajectory.atom_types atom_type_names = [] @@ -243,6 +288,14 @@ def run(self) -> WFData: # frame_data should always be shaped: n_probes,nkx,nky,n_layers,1 (idk why there's a trailing 1) cache_exists,frame_data = checkCache(cache_file,self.cache_levels) + if not os.path.exists(self.output_dir / f"kx.npy"): + np.save(self.output_dir / f"kx.npy",to_cpu(self.kxs)) + np.save(self.output_dir / f"ky.npy",to_cpu(self.kys)) + if len(self.kxs)!=len(self.kxs_uncrop) and not os.path.exists(self.output_dir / f"kx_uncrop.npy"): + np.save(self.output_dir / f"kx_uncrop.npy",to_cpu(self.kxs_uncrop)) + if len(self.kys)!=len(self.kys_uncrop) and not os.path.exists(self.output_dir / f"ky_uncrop.npy"): + np.save(self.output_dir / f"ky_uncrop.npy",to_cpu(self.kys_uncrop)) + if cache_exists: frames_cached += 1 else: @@ -335,7 +388,7 @@ def run(self) -> WFData: # Handle cleanup if self.cleanup_temp_files: logger.info("Cleaning up cache files...") - for frame_idx in range(n_frames): + for frame_idx in range(self.n_frames): cache_file = self.output_dir / f"frame_{frame_idx}.npy" if cache_file.exists(): cache_file.unlink() diff --git a/src/pyslice/multislice/multislice.py b/src/pyslice/multislice/multislice.py index 7599115..0f2f8ff 100644 --- a/src/pyslice/multislice/multislice.py +++ b/src/pyslice/multislice/multislice.py @@ -1,7 +1,7 @@ import numpy as np from tqdm import tqdm import logging -from ..backend import zeros,mean,ones +from ..backend import zeros,mean,ones,to_cpu try: import torch ; xp = torch @@ -29,15 +29,48 @@ logger = logging.getLogger(__name__) +def antialias_aperture(kxs, kys, cutoff_fraction=2/3, taper_width=0.02): + """Compute a 2/3 Nyquist anti-aliasing aperture in k-space. + + Bandwidth-limits to ``cutoff_fraction`` of k_max with a smooth cosine + taper to avoid ringing from a hard cutoff. Returns a 2D real-valued + array (1 inside, tapers to 0 outside). + """ + kx_max = xp.amax(xp.abs(kxs)) + ky_max = xp.amax(xp.abs(kys)) + k_max = min(float(kx_max), float(ky_max)) # Nyquist = 1/(2*sampling) + k_cutoff = cutoff_fraction * k_max + + kx_grid, ky_grid = xp.meshgrid(kxs, kys, indexing='ij') + k_r = xp.sqrt(kx_grid**2 + ky_grid**2) + + # Smooth cosine taper from (k_cutoff - taper) to k_cutoff, strictly zero above k_cutoff + taper = taper_width * k_max + aperture = xp.ones_like(k_r) + mask_taper = (k_r > k_cutoff - taper) & (k_r < k_cutoff) + mask_outer = k_r >= k_cutoff + if TORCH_AVAILABLE and hasattr(k_r, 'device'): + aperture[mask_taper] = 0.5 * (1 + torch.cos(torch.pi * (k_r[mask_taper] - k_cutoff + taper) / taper)) + else: + aperture[mask_taper] = 0.5 * (1 + np.cos(np.pi * (k_r[mask_taper] - k_cutoff + taper) / taper)) + aperture[mask_outer] = 0.0 + return aperture + m_electron = 9.109383e-31 # mass of an electron, kg -q_electron = 1.602177e-19 # charge of an electron, J / eV or kg m^2/s^2 / eV +q_electron = 1.602177e-19 # charge of an electron, J / eV or kg m^2/s^2 / eV c_light = 299792458.0 # speed of light, m / s h_planck = 6.62607015e-34 # m^2 kg / s +#what if we use units of Å throughout, instead of m? +#m_electron = 9.109383e-19 # mass of an electron, nano-gram (ng) +#q_electron = 1.602177e1 # charge of an electron, kg Å^2/s^2 / eV +#c_light = 2.99792458e18 # speed of light, Å / s +#h_planck = 6.62607015e-14 # Å^2 ng / s def m_effective(eV): """Relativistic correction: E=m*c^2, so m=E/c^2, in kg""" return m_electron + eV * q_electron / c_light**2 + # units [ kg ] [ eV ] [ kg m² s⁻² eV⁻¹ ] [ m⁻² s² ] def wavelength(eV): """ @@ -243,7 +276,7 @@ def plot(self,filename=None,title=None): import matplotlib.pyplot as plt fig, ax = plt.subplots() # calling self.array should convert to CPU/numpy - array = np.mean(np.absolute(self.array[:,:,:,:]),axis=1)[0,:,:] # positional,summable,x,y indices + array = np.mean(np.absolute(self.array[:,:,::-1,:]),axis=1)[0,:,:] # positional,summable,x,y indices array=array.T # imshow convention: y,x. our convention: x,y plot_array = np.absolute(array)**.25 @@ -340,9 +373,9 @@ def applyShifts(self): # Create shifted probe using phase ramp in k-space probe_k = xp.fft.fft2(self._array[:,i,:,:]) # positional,summable,x,y - # Apply phase ramp for spatial shift - kx_shift = xp.exp(2j * xp.pi * self.kxs[None,:, None] * (px-self.lx/2) ) - ky_shift = xp.exp(2j * xp.pi * self.kys[None,None, :] * (py-self.ly/2) ) + # Apply phase ramp for spatial shift (negative sign = shift right) + kx_shift = xp.exp(-2j * xp.pi * self.kxs[None,:, None] * (px-self.lx/2) ) + ky_shift = xp.exp(-2j * xp.pi * self.kys[None,None, :] * (py-self.ly/2) ) probe_k_shifted = probe_k * kx_shift * ky_shift # Convert back to real space @@ -446,9 +479,9 @@ def create_batched_probes(base_probe, probe_positions, device=None): # Create shifted probe using phase ramp in k-space probe_k = xp.fft.fft2(base_probe._array) - # Apply phase ramp for spatial shift - kx_shift = xp.exp(2j * xp.pi * base_probe.kxs[:, None] * (px-lx/2) ) - ky_shift = xp.exp(2j * xp.pi * base_probe.kys[None, :] * (py-ly/2) ) + # Apply phase ramp for spatial shift (negative sign = shift right) + kx_shift = xp.exp(-2j * xp.pi * base_probe.kxs[:, None] * (px-lx/2) ) + ky_shift = xp.exp(-2j * xp.pi * base_probe.kys[None, :] * (py-ly/2) ) probe_k_shifted = probe_k * kx_shift * ky_shift # Convert back to real space @@ -463,6 +496,12 @@ def create_batched_probes(base_probe, probe_positions, device=None): return Probe(base_probe.xs, base_probe.ys, base_probe.mrad, base_probe.eV, array=array, device=base_probe.device) +# Given a real-space entrance wave, and a potential (or object), calculate the exit wave: ψ₁ -> O -> ψ₂ +# From Kirkland2010: +# propagator P = exp(-i π λ dz q²), Eq 6.65 +# transmission function t = exp(i σ O) where O is our object (or potential slice), Eq 6.59 +# ψ₂ = ℱ⁻¹[ P * ℱ[ t * ψ₁ ] ], Eq 6.67, noting the relationship: 𝒞[ f(x),g(x) ] = ℱ⁻¹[ ℱ[f(x)] * ℱ[g(x)] ] = ℱ⁻¹[ f(k) * g(k) ] +# or as code: array = t * array ; fft_array = fft(array) ; propagated_fft = P * fft_array ; array = ifft(propagated_fft) def Propagate(probe, potential, device=None, progress=False, onthefly=True, store_all_slices=False): """ PyTorch-accelerated multislice propagation function. @@ -485,7 +524,6 @@ def Propagate(probe, potential, device=None, progress=False, onthefly=True, stor raise ImportError("PyTorch not available. Please install PyTorch.") # Initialize wavefunction with probe(s) - shape: (n_probes, nx, ny) - #print(probe._array.shape) nc,npt,nx,ny = probe._array.shape #; print("nc,npt,nx,ny",nc,npt,nx,ny) array = probe._array.reshape((nc*npt,nx,ny)) # "flatten" first two indices probe_wavelengths = probe.wavelengths[:,None]*ones(npt)[None,:] # also expand wavelengths and eVs arrays to cover all probe positions npt @@ -497,7 +535,7 @@ def Propagate(probe, potential, device=None, progress=False, onthefly=True, stor E0_eV = m_electron * c_light**2 / q_electron sigma = (2 * np.pi) / (probe_wavelengths * probe_eVs) * \ (E0_eV + probe_eVs) / (2 * E0_eV + probe_eVs) # wavelength and eVs now have length of n_probes - + #print("propagate sigma",sigma) #if TORCH_AVAILABLE: # #sigma_dtype = torch.float32 if device.type == 'mps' else torch.float64 # sigma = torch.tensor(sigma, dtype=float_dtype, device=device) @@ -509,7 +547,12 @@ def Propagate(probe, potential, device=None, progress=False, onthefly=True, stor # All tensors should already be on the correct device from creation kx_grid, ky_grid = xp.meshgrid(potential.kxs, potential.kys, indexing='ij') k_squared = kx_grid**2 + ky_grid**2 - P = xp.exp(-1j * xp.pi * probe_wavelengths[:,None,None] * dz * k_squared[None,:,:]) + + # Precompute 2/3 Nyquist anti-aliasing aperture for bandwidth-limiting transmission functions + aa_aperture = antialias_aperture(potential.kxs, potential.kys) + + # Fold anti-aliasing aperture into propagator to bandwidth-limit wavefunction at every slice + P = xp.exp(-1j * xp.pi * probe_wavelengths[:,None,None] * dz * k_squared[None,:,:]) * aa_aperture[None,:,:] # Kirkland2010 Eq 6.65 if progress: localtqdm = tqdm @@ -532,10 +575,14 @@ def localtqdm(iterator): potential_slice = potential.calculateSlice(z) else: potential_slice = potential._array[:, :, z] - t = xp.exp(1j * sigma[:,None,None] * potential_slice[None,:,:]) + t = xp.exp(1j * sigma[:,None,None] * potential_slice[None,:,:]) # Kirkland2010 Eq 6.59 + + # Bandwidth-limit the transmission function to 2/3 Nyquist to prevent aliasing + kwarg = {"dim":(-2,-1)} if TORCH_AVAILABLE else {"axes":(-2,-1)} + t = xp.fft.ifft2(xp.fft.fft2(t, **kwarg) * aa_aperture[None,:,:], **kwarg) # Apply transmission to all probes: ψ' = t × ψ - # Broadcasting: t[nx,ny] * array[n_probes,nx,ny] = array[n_probes,nx,ny] + # Broadcasting: t[n_probes,nx,ny] * array[n_probes,nx,ny] = array[n_probes,nx,ny] array = t * array # Store wavefunction at this slice if requested (after transmission) @@ -570,3 +617,68 @@ def localtqdm(iterator): # return array.squeeze(0) return array # okay for Propagate to return a Tensor. we probably don't want to move things off-gpu yet +# Given a real-space entrance and real-space exit wave, calculate the object the wave must have passed through: ψ₁ -> O -> ψ₂, given ψ₁,ψ₂, find O +# From Kirkland2010: +# propagator P = exp(-i π λ dz q²), Eq 6.65 +# transmission function t = exp(i σ O) where O is our object (or potential slice), Eq 6.59 +# ψ₂ = ℱ⁻¹[ P * ℱ[ t * ψ₁ ] ], Eq 6.67, noting the relationship: 𝒞[ f(x),g(x) ] = ℱ⁻¹[ ℱ[f(x)] * ℱ[g(x)] ] = ℱ⁻¹[ f(k) * g(k) ] +# or as code: array = t * array ; fft_array = fft(array) ; propagated_fft = P * fft_array ; array = ifft(propagated_fft) +# SO, to determine O from ψ₁ and ψ₂: +# ℱ[ ψ₂ ] = P * ℱ[ t * ψ₁ ] +# ℱ[ ψ₂ ]/P = ℱ[ t * ψ₁ ] +# ℱ⁻¹[ ℱ[ ψ₂ ]/P ] = t * ψ₁ +# t = ℱ⁻¹[ ℱ[ ψ₂ ]/P ]/ψ₁ +# exp(i σ O) = ℱ⁻¹[ ℱ[ ψ₂ ]/P ]/ψ₁ +# i σ O = log( ℱ⁻¹[ ℱ[ ψ₂ ]/P ]/ψ₁ ) +# O = log( ℱ⁻¹[ ℱ[ ψ₂ ]/P ]/ψ₁ ) / i / σ +def calculateObject(probe,exitwave,guessedObject,weighting=.5,dz=0.5,damping=.01): + + import matplotlib.pyplot as plt + #fig, axs = plt.subplots(1,2) + #axs[0].imshow(to_cpu(xp.absolute(exitwave)), cmap="inferno") ; axs[0].set_title("exit wave") + #axs[1].imshow(to_cpu(xp.absolute(probe._array[0,0,:,:])), cmap="inferno") ; axs[1].set_title("entrance wave") + #plt.show() + + nc,npt,nx,ny = probe._array.shape #; print("nc,npt,nx,ny",nc,npt,nx,ny) + psi1 = probe._array[0,0,:,:] # select first probe array only for now! + lamda = probe.wavelengths[0] + eV = probe.eVs[0] + + # Calculate interaction parameter (Kirkland Eq 5.6) + E0_eV = m_electron * c_light**2 / q_electron + sigma = (2 * np.pi) / (lamda * eV) * \ + (E0_eV + eV) / (2 * E0_eV + eV) # wavelength and eVs now have length of n_probes + + # Pre-compute propagation operator in k-space (Fresnel propagation) + # All tensors should already be on the correct device from creation + kx_grid, ky_grid = xp.meshgrid(probe.kxs, probe.kys, indexing='ij') # TODO use probe.kxs instead, to free ourselves of the need to pass a potential + k_squared = kx_grid**2 + ky_grid**2 + + P = xp.exp(-1j * xp.pi * lamda * dz * k_squared[:,:]) # P in Kirkland2010 Eq 6.65 is exp(-i...), so to divide by P, we can use Pp, exp(+j...) + + # t = ℱ⁻¹[ ℱ[ ψ₂ ]/P ]/ψ₁ + t = xp.fft.ifft2(xp.fft.fft2(exitwave)/P)/psi1 + # t = exp(i σ O) --> O = log(t)/i/σ + O = xp.log(t)/1j/sigma + # WHAT IS exp(i σ O) REALLY DOING? exp(iϕ) is a sinusoid. an only-real object is applying a phase shift? can we do calculate an angle instead? + O = np.angle(t)/sigma + # consider: + # instead of division, should i multiply by complex conjugate? (not technically the same, but it will deal with near-zeros) + # instead of log, should i use: https://en.wikipedia.org/wiki/Complex_logarithm#Calculating_the_principal_value + # should I apply a probe amplitude masking function? + + #fig, axs = plt.subplots(1,3) + #axs[0].imshow(to_cpu(xp.absolute(exitwave)), cmap="inferno") ; axs[0].set_title("exit wave") + #axs[1].imshow(to_cpu(xp.absolute(probe._array[0,0,:,:])), cmap="inferno") ; axs[1].set_title("entrance wave") + #axs[2].imshow(to_cpu(xp.absolute(O)**.1), cmap="inferno") ; axs[2].set_title("object") + #plt.show() + + delta=(O-guessedObject) + + # probe amplitude masking function: zeros-out points where probe intensity is zero, without giving probe features as features in your potential + delta*=xp.absolute(psi1)/(xp.absolute(psi1)+damping*xp.amax(xp.absolute(psi1))) + + return delta*weighting + + + diff --git a/src/pyslice/multislice/potentials.py b/src/pyslice/multislice/potentials.py index 48e8134..092e7ac 100644 --- a/src/pyslice/multislice/potentials.py +++ b/src/pyslice/multislice/potentials.py @@ -2,6 +2,7 @@ from pathlib import Path import logging,os from tqdm import tqdm +from ..backend import to_cpu,mean try: import torch ; xp = torch @@ -186,8 +187,8 @@ def loadKirkland(device='cpu'): else: kirklandABCDs = np.asarray(kirkland_params) -class Potential: - def __init__(self, xs, ys, zs, positions, atomTypes, kind="kirkland", device=None, slice_axis=2, progress=False, cache_dir=None, frame_idx=None): +class Potential: + def __init__(self, xs, ys, zs, positions=None, atomTypes=None, array=None, kind="kirkland", device=None, slice_axis=2, progress=False, cache_dir=None, frame_idx=None): # Set up device and backend first if TORCH_AVAILABLE: # Auto-detect device if not specified @@ -212,7 +213,12 @@ def __init__(self, xs, ys, zs, positions, atomTypes, kind="kirkland", device=Non self.xs = torch.tensor(xs, dtype=self.dtype, device=device) self.ys = torch.tensor(ys, dtype=self.dtype, device=device) self.zs = torch.tensor(zs, dtype=self.dtype, device=device) - positions = torch.tensor(positions, dtype=self.dtype, device=device) + if positions is not None: + positions = torch.tensor(positions, dtype=self.dtype, device=device) + if array is not None: + self._array = torch.tensor(array, device=device) + else: + self._array = None else: if device is not None: raise ImportError("PyTorch not available. Please install PyTorch.") @@ -223,6 +229,7 @@ def __init__(self, xs, ys, zs, positions, atomTypes, kind="kirkland", device=Non self.xs = xs self.ys = ys self.zs = zs + self._array = array self.nx = len(xs) self.ny = len(ys) @@ -255,32 +262,35 @@ def __init__(self, xs, ys, zs, positions, atomTypes, kind="kirkland", device=Non qsq = self.kxs[:, None]**2 + self.kys[None, :]**2 # Convert atom types to atomic numbers if needed - unique_atom_types = set(atomTypes) - atomic_numbers = [] - for at in atomTypes: - if isinstance(at, str): - atomic_numbers.append(getZfromElementName(at)) - else: - atomic_numbers.append(at) - if TORCH_AVAILABLE: - atomic_numbers = torch.tensor(atomic_numbers, device=device) - - # OPTIMIZATION 1: Compute form factors once per atom type on GPU - form_factors = {} - for at in unique_atom_types: - if kind == "kirkland": + if atomTypes is not None: + unique_atom_types = set(atomTypes) + atomic_numbers = [] + for at in atomTypes: if isinstance(at, str): - Z = getZfromElementName(at) + atomic_numbers.append(getZfromElementName(at)) else: - Z = at - form_factors[at] = kirkland(qsq, Z) - elif kind == "gauss": - form_factors[at] = torch.exp(-1**2 * qsq / 2) + atomic_numbers.append(at) + if TORCH_AVAILABLE: + atomic_numbers = torch.tensor(atomic_numbers, device=device) + + # OPTIMIZATION 1: Compute form factors once per atom type on GPU + form_factors = {} + for at in unique_atom_types: + if kind == "kirkland": + if isinstance(at, str): + Z = getZfromElementName(at) + else: + Z = at + form_factors[at] = kirkland(qsq, Z) + elif kind == "gauss": + form_factors[at] = torch.exp(-1**2 * qsq / 2) self.cache_dir = cache_dir self.frame_idx = frame_idx def calculateSlice(slice_idx): + if self._array is not None: + return self._array[:,:,slice_idx] # check for caching cache_file = None if self.cache_dir is not None: @@ -367,10 +377,11 @@ def calculateSlice(slice_idx): return Z self.calculateSlice = calculateSlice - self._array = None + #self._array = None def build(self,progress=False): - + if self._array is not None: + return # Initialize potential array using xp with conditional device device_kwargs = {'device': self.device } if self.use_torch else {} potential_real = xp.zeros((self.nx, self.ny, self.n_slices), dtype=float_dtype, **device_kwargs) @@ -389,17 +400,15 @@ def localtqdm(iterator): # Store tensor version for potential GPU operations self._array = potential_real + def flatten(self): + self._array = mean(self._array,axis=2,keepdims=True) + self.nz = 1 + self.zs=self.zs[:1] + @property def array(self): - return self.to_cpu() - - def to_cpu(self): - """Convert tensors back to CPU NumPy arrays.""" - if self.use_torch: - return self._array.cpu().numpy() - else: - return self._array - + return to_cpu(self._array) + def to_device(self, device): """Move tensor data to specified device.""" if hasattr(self, 'array_torch'): @@ -413,7 +422,7 @@ def plot(self,filename=None): import matplotlib.pyplot as plt fig, ax = plt.subplots() - array = np.sum(np.absolute(self.array),axis=2).T # imshow convention: y,x. our convention: x,y + array = np.sum(np.absolute(self.array),axis=2).T[::-1,:] # imshow convention: y,x. our convention: x,y, and flip y (0,0 upper-left) # Convert to CPU if on GPU/MPS device #if hasattr(array, 'cpu'): # array = array.cpu() diff --git a/src/pyslice/multislice/trajectory.py b/src/pyslice/multislice/trajectory.py index f778eef..3039fb6 100644 --- a/src/pyslice/multislice/trajectory.py +++ b/src/pyslice/multislice/trajectory.py @@ -4,6 +4,7 @@ from dataclasses import dataclass import numpy as np from typing import List, Tuple, Optional +from ase import Atoms @dataclass class Trajectory: @@ -329,7 +330,7 @@ def generate_random_displacements(self,n_displacements,sigma,seed=None): na=len(self.positions[0]) if seed is not None: np.random.seed(seed) - dxyz=np.random.random(size=(n_displacements,na,3))*sigma + dxyz=np.random.normal(0,sigma,size=(n_displacements,na,3)) positions = self.positions[0]+dxyz return Trajectory( @@ -503,3 +504,8 @@ def rotate_to(self, direction: Tuple[int, int, int]) -> 'Trajectory': box_matrix=rotated_box_matrix, timestep=self.timestep ) + + # returns an ase object + def to_ase(self): + return Atoms(''.join(self.atom_types), positions=self.positions[0], cell=np.diag(self.box_matrix),pbc=True) + diff --git a/src/pyslice/postprocessing/haadf_data.py b/src/pyslice/postprocessing/haadf_data.py index 671f544..74f9c78 100644 --- a/src/pyslice/postprocessing/haadf_data.py +++ b/src/pyslice/postprocessing/haadf_data.py @@ -6,8 +6,10 @@ from pathlib import Path import logging from .wf_data import WFData -from ..data import Signal, Dimensions, Dimension, GeneralMetadata -from ..backend import zeros +from ..data.pyslice_serial import PySliceSerial, Signal, Dimensions, Dimension, Metadata +#from ..data import Signal, Dimensions, Dimension, GeneralMetadata +from pyslice.backend import zeros,to_cpu +#from ..data.pyslice_serial import PySliceSerial logger = logging.getLogger(__name__) @@ -34,7 +36,7 @@ float_dtype = np.float64 -class HAADFData(Signal): +class HAADFData(PySliceSerial, Signal): """ Data structure for HAADF (High Angle Annular Dark Field) imaging data. @@ -49,6 +51,14 @@ class HAADFData(Signal): cache_dir: Path to cache directory. """ + _sea_config = { + 'tensor_attrs': ['_kxs', '_kys', '_xs', '_ys', '_array', 'data'], + 'path_attrs': ['cache_dir'], + 'tuple_list_attrs': ['probe_positions'], + 'exclude_attrs': ['probe', '_wf_array'], + 'force_datasets': ['_array', 'probe_positions', '_kxs', '_kys', '_xs', '_ys'], + } + def __init__(self, wf_data: WFData) -> None: """ Initialize HAADFData from WFData. @@ -71,35 +81,28 @@ def __init__(self, wf_data: WFData) -> None: self._xs = wf_data.probe_xs self._ys = wf_data.probe_ys - # Build placeholder dimensions (will be updated after calculateADF) - dimensions = Dimensions([ - Dimension(name='x', space='position', units='Å', values=np.array([0])), - Dimension(name='y', space='position', units='Å', values=np.array([0])), - ], nav_dimensions=[0, 1], sig_dimensions=[]) - - # Build metadata - metadata_dict = { - 'General': { - 'title': 'HAADF Image', - 'signal_type': 'HAADF' - }, - 'Simulation': { - 'voltage_eV': float(self.probe.eV), - 'wavelength_A': float(self.probe.wavelength), - 'aperture_mrad': float(self.probe.mrad), - 'probe_positions': [list(p) for p in self.probe_positions], + if Dimensions is not None: + # Build placeholder dimensions (will be updated after calculateADF) + self.dimensions = Dimensions([ + Dimension(name='x', space='position', units='Å', values=np.array([0])), + Dimension(name='y', space='position', units='Å', values=np.array([0])), + ], nav_dimensions=[0, 1], sig_dimensions=[]) + + # Build metadata + metadata_dict = { + 'General': { + 'title': 'HAADF Image', + 'signal_type': 'HAADF' + }, + 'Simulation': { + 'voltage_eV': float(self.probe.eV), + 'wavelength_A': float(self.probe.wavelength), + 'aperture_mrad': float(self.probe.mrad), + 'probe_positions': [list(p) for p in self.probe_positions], + } } - } - metadata = GeneralMetadata(metadata_dict) - - # Initialize Signal base class - super().__init__( - data=None, # We'll override the data property - name='HAADFData', - dimensions=dimensions, - signal_type='Image', - metadata=metadata - ) + self.metadata = Metadata(metadata_dict) + self.sea_type="Signal" @property def data(self): @@ -126,7 +129,7 @@ def adf(self, value): @property def array(self): """Alias for adf (backward compatibility).""" - return self._array + return to_cpu(self._array) def __getattr__(self, name): """Auto-convert coordinate arrays from tensor to numpy on access.""" @@ -195,7 +198,7 @@ def calculateADF(self, inner_mrad: float = 45, outer_mrad: float = 150, preview: plt.show() collected = self._wf_array * mask[None,None,None,:,:,None] # c,x,y,t,kx,ky,l indices, mask is kx,ky - self._array = xp.mean(xp.sum(xp.absolute(collected),axis=(4,5)),axis=(0,3,4)) # sum over kx,ky, mean over c,t,l + self._array = xp.mean(xp.sum(xp.absolute(collected)**2,axis=(4,5)),axis=(0,3,4)) # sum |ψ|² over kx,ky, mean over c,t,l # Update dimensions with computed xs, ys def to_numpy(x): @@ -206,13 +209,14 @@ def to_numpy(x): xs_np = to_numpy(self._xs) ys_np = to_numpy(self._ys) - self._local_dimensions = Dimensions([ - Dimension(name='x', space='position', units='Å', values=xs_np), - Dimension(name='y', space='position', units='Å', values=ys_np), - ], nav_dimensions=[0, 1], sig_dimensions=[]) + if Dimensions is not None: + self._local_dimensions = Dimensions([ + Dimension(name='x', space='position', units='Å', values=xs_np), + Dimension(name='y', space='position', units='Å', values=ys_np), + ], nav_dimensions=[0, 1], sig_dimensions=[]) - # Update metadata with detector settings - if hasattr(self.metadata, 'Simulation'): + # Update metadata with detector settings + #if hasattr(self.signal.metadata, 'Simulation'): self.metadata.Simulation.inner_mrad = inner_mrad self.metadata.Simulation.outer_mrad = outer_mrad @@ -231,16 +235,9 @@ def plot(self, filename=None): raise RuntimeError("calculateADF() must be called before plotting") fig, ax = plt.subplots() - array = self._array.T # imshow convention: y,x. our convention: x,y - - if TORCH_AVAILABLE and hasattr(array, 'cpu'): - array = array.cpu().numpy() - if TORCH_AVAILABLE and hasattr(self._xs, 'cpu'): - xs = self._xs.cpu().numpy() - ys = self._ys.cpu().numpy() - else: - xs = np.asarray(self._xs) - ys = np.asarray(self._ys) + array = self.array.T[::-1,:] # imshow convention: y,x. our convention: x,y, and flip y (0,0 upper-left) + xs = to_cpu(self._xs) + ys = to_cpu(self._ys) extent = (np.amin(xs), np.amax(xs), np.amin(ys), np.amax(ys)) ax.imshow(array, cmap="inferno", extent=extent) diff --git a/src/pyslice/postprocessing/tacaw_data.py b/src/pyslice/postprocessing/tacaw_data.py index ba4e0df..c49c0d1 100644 --- a/src/pyslice/postprocessing/tacaw_data.py +++ b/src/pyslice/postprocessing/tacaw_data.py @@ -6,7 +6,9 @@ from pathlib import Path import logging, os from .wf_data import WFData -from ..data import Signal, Dimensions, Dimension, GeneralMetadata +from ..data.pyslice_serial import PySliceSerial, Signal, Dimensions, Dimension, Metadata +#from ..data import Signal, Dimensions, Dimension, GeneralMetadata +#from ..data.pyslice_serial import PySliceSerial from pyslice.backend import to_cpu,fft,fftshift,mean from tqdm import tqdm @@ -35,7 +37,7 @@ float_dtype = np.float64 -class TACAWData(Signal): +class TACAWData(PySliceSerial, Signal): """ Data structure for storing TACAW EELS results with format: probe_positions, frequency, kx, ky. @@ -53,6 +55,14 @@ class TACAWData(Signal): cache_dir: Path to cache directory. """ + _sea_config = { + 'tensor_attrs': ['_kxs', '_kys', '_xs', '_ys', '_time', '_layer', '_frequencies', '_array', 'data'], + 'path_attrs': ['cache_dir'], + 'tuple_list_attrs': ['probe_positions'], + 'exclude_attrs': ['probe', '_wf_array'], + 'force_datasets': ['_array', 'probe_positions', '_kxs', '_kys', '_xs', '_ys', '_time', '_layer', '_frequencies'], + } + def __init__(self, wf_data: WFData, layer_index: int = None, keep_complex: bool = False, chunkFFT: bool = False) -> None: """ Initialize TACAWData from WFData by performing FFT. @@ -77,7 +87,7 @@ def __init__(self, wf_data: WFData, layer_index: int = None, keep_complex: bool # Store reference to source WFData array for FFT computation self._wf_array = wf_data._array - print("tacaw > wfdata",wf_data._array.shape) + #print("tacaw > wfdata",wf_data._array.shape) # Initialize intensity as None, will be set by fft_from_wf_data self._array = None @@ -101,40 +111,34 @@ def to_numpy(x): kxs_arr = to_numpy(self._kxs) kys_arr = to_numpy(self._kys) - dimensions = Dimensions([ - Dimension(name='probe', space='position', - values=np.arange(len(self.probe_positions))), - Dimension(name='frequency', space='spectral', units='THz', - values=freq_arr), - Dimension(name='kx', space='scattering', units='Å⁻¹', - values=kxs_arr), - Dimension(name='ky', space='scattering', units='Å⁻¹', - values=kys_arr), - ], nav_dimensions=[0, 1], sig_dimensions=[2, 3]) - - # Build metadata - metadata_dict = { - 'General': { - 'title': 'TACAW Intensity', - 'signal_type': 'TACAW' - }, - 'Simulation': { - 'voltage_eV': float(self.probe.eV), - 'wavelength_A': float(self.probe.wavelength), - 'aperture_mrad': float(self.probe.mrad), - 'probe_positions': [list(p) for p in self.probe_positions], + if Dimensions is not None: + + self.dimensions = Dimensions([ + Dimension(name='probe', space='position', + values=np.arange(len(self.probe_positions))), + Dimension(name='frequency', space='spectral', units='THz', + values=freq_arr), + Dimension(name='kx', space='scattering', units='Å⁻¹', + values=kxs_arr), + Dimension(name='ky', space='scattering', units='Å⁻¹', + values=kys_arr), + ], nav_dimensions=[0, 1], sig_dimensions=[2, 3]) + + # Build metadata + metadata_dict = { + 'General': { + 'title': 'TACAW Intensity', + 'signal_type': 'TACAW' + }, + 'Simulation': { + 'voltage_eV': float(self.probe.eV), + 'wavelength_A': float(self.probe.wavelength), + 'aperture_mrad': float(self.probe.mrad), + 'probe_positions': [list(p) for p in self.probe_positions], + } } - } - metadata = GeneralMetadata(metadata_dict) - - # Initialize Signal base class (this will set self.data = None) - super().__init__( - data=None, - name='TACAWData', - dimensions=dimensions, - signal_type='2D-EELS', - metadata=metadata - ) + self.metadata = Metadata(metadata_dict) + self.sea_type="Signal" # Restore computed values AFTER super().__init__ self._array = computed_array @@ -157,9 +161,7 @@ def data(self): """Lazy conversion to numpy for Signal compatibility.""" if self._array is None: return None - if hasattr(self._array, 'cpu'): - return self._array.cpu().numpy() - return np.asarray(self._array) + return to_cpu(self._array) @data.setter def data(self, value): @@ -177,7 +179,7 @@ def intensity(self, value): @property def array(self): """Alias for intensity (backward compatibility with WFData interface).""" - return self._array + return to_cpu(self._array) def _fft_from_wf_data(self, layer_index: int = None): """ @@ -267,6 +269,8 @@ def _fft_from_wf_data(self, layer_index: int = None): # else: # self._array = np.abs(wf_fft)**2 + # Ensure cache directory exists (may have been cleaned up by calculator) + self.cache_dir.mkdir(parents=True, exist_ok=True) np.save(self.cache_dir / "tacaw_freq.npy", self._frequencies) np.save(self.cache_dir / "tacaw.npy", self._array.detach().cpu().numpy() if TORCH_AVAILABLE and hasattr(self._array, 'cpu') else self._array) @@ -285,6 +289,7 @@ def spectrum(self, probe_index: int = None) -> np.ndarray: Returns: Spectrum array (frequency intensity) """ + if probe_index is None: # Average over all probe positions all_spectra = [] @@ -427,7 +432,7 @@ def spectral_diffraction(self, frequency: float, probe_index: int = None, space: return spectral_diffraction - def masked_spectrum(self, mask: np.ndarray, probe_index: int = None) -> np.ndarray: + def masked_spectrum(self, mask: np.ndarray|dict|None = None, probe_index: int = None, preview=False) -> np.ndarray: """ Extract spectrum with spatial masking in k-space. @@ -441,38 +446,38 @@ def masked_spectrum(self, mask: np.ndarray, probe_index: int = None) -> np.ndarr kxs = to_cpu(self.kxs) kys = to_cpu(self.kys) - if mask.shape != (len(kxs), len(kys)): + if mask is None: + mask = np.zeros((len(kxs),len(kys)))+1 + elif isinstance(mask,dict): + cx,cy=mask.get("center",(0,0)) + if mask["shape"] == "round": + r=mask["radius"] + radii = np.sqrt((kxs[:,None]-cx)**2+(kys[None,:]-cy)**2) + mask = np.zeros((len(kxs),len(kys))) + mask[radii<=r]=1 + + elif mask.shape != (len(kxs), len(kys)): raise ValueError(f"Mask shape {mask.shape} doesn't match k-space shape ({len(kxs)}, {len(kys)})") if probe_index is None: - # Average over all probe positions - all_masked_spectra = [] - for i in range(len(self.probe_positions)): - probe_intensity = self._array[i] # Shape: (frequency, kx, ky) - masked_intensity = probe_intensity * mask[None, :, :] # Broadcast mask to all frequencies - masked_spectrum = xp.sum(masked_intensity, axis=(1, 2)) # Sum over masked k-space - all_masked_spectra.append(masked_spectrum) - - # Average all masked spectra - if TORCH_AVAILABLE and hasattr(all_masked_spectra[0], 'cpu'): - all_masked_spectra = [ms.cpu().numpy() for ms in all_masked_spectra] - masked_spectrum = np.mean(all_masked_spectra, axis=0) - else: - if probe_index >= len(self.probe_positions): - raise ValueError(f"Probe index {probe_index} out of range") - - # Extract intensity data for this probe - probe_intensity = self._array[probe_index] # Shape: (frequency, kx, ky) - - # Apply spatial mask in k-space - masked_intensity = probe_intensity * mask[None, :, :] # Broadcast mask to all frequencies - masked_spectrum = xp.sum(masked_intensity, axis=(1, 2)) # Sum over masked k-space - - # Convert to numpy if PyTorch tensor - if TORCH_AVAILABLE and hasattr(masked_spectrum, 'cpu'): - masked_spectrum = masked_spectrum.cpu().numpy() - - return masked_spectrum + probe_index = np.arange(len(self.probe_positions)) + elif isinstance(probe_index,int): + probe_index=[probe_index] + spectra = [] + for i in probe_index: + masked = self._array[i] * mask[None,:,:] + if preview: + import matplotlib.pyplot as plt + fig, ax = plt.subplots() + extent = ( np.amin(kxs) , np.amax(kxs) , np.amin(kys) , np.amax(kys) ) + ax.imshow(to_cpu(xp.sum(masked,axis=0).T)[::-1,:], cmap="inferno",extent=extent,aspect=1) + ax.set_xlabel("kx") + ax.set_ylabel("ky") + ax.set_title("masked_spectrum - preview") + plt.show() + preview=False + spectra.append(to_cpu(xp.sum(masked,axis=(1,2)))) + return np.mean(spectra,axis=0) def dispersion(self, kx_path: np.ndarray, ky_path: np.ndarray, probe_index: int = None, space: str = "reciprocal") -> np.ndarray: """ @@ -579,7 +584,6 @@ def plot(self,intensities,xvals,yvals,xlabel="kx ($\\AA^{-1}$)",ylabel="ky ($\\A else: plt.show() - class SEDData(TACAWData): """ Data structure for SED (Spectral Energy Density) calculations. diff --git a/src/pyslice/postprocessing/wf_data.py b/src/pyslice/postprocessing/wf_data.py index 7d0dcaf..128fef1 100644 --- a/src/pyslice/postprocessing/wf_data.py +++ b/src/pyslice/postprocessing/wf_data.py @@ -4,7 +4,7 @@ import numpy as np from typing import List, Tuple, Optional from ..multislice.multislice import Probe,aberrationFunction -from ..data import Signal, Dimensions, Dimension, GeneralMetadata +from ..data.pyslice_serial import PySliceSerial, Signal, Dimensions, Dimension, Metadata from pathlib import Path from ..backend import mean,ones,zeros @@ -31,7 +31,7 @@ float_dtype = np.float64 -class WFData(Signal): +class WFData(PySliceSerial, Signal): """ Data structure for wave function data with format: probe_positions, frame, kx, ky, layer. @@ -50,6 +50,14 @@ class WFData(Signal): cache_dir: Path to cache directory. """ + _sea_config = { + 'tensor_attrs': ['_kxs', '_kys', '_xs', '_ys', '_time', '_layer', '_array'], + 'path_attrs': ['cache_dir'], + 'tuple_list_attrs': ['probe_positions'], + 'exclude_attrs': ['probe'], + 'force_datasets': ['_array', 'probe_positions', '_kxs', '_kys', '_xs', '_ys', '_time', '_layer'], + } + def __init__( self, probe_positions: List[Tuple[float, float]], @@ -90,46 +98,37 @@ def to_numpy(x): kys_arr = to_numpy(kys) layer_arr = to_numpy(layer) if layer is not None else np.array([0]) - dimensions = Dimensions([ - Dimension(name='probe', space='position', - values=np.arange(len(probe_positions))), - Dimension(name='time', space='temporal', units='ps', - values=time_arr), - Dimension(name='kx', space='scattering', units='Å⁻¹', - values=kxs_arr), - Dimension(name='ky', space='scattering', units='Å⁻¹', - values=kys_arr), - Dimension(name='layer', space='position', - values=layer_arr), - ], nav_dimensions=[0, 1], sig_dimensions=[2, 3, 4]) - - # Build metadata from simulation parameters - # Flatten probe_positions for HDF5 compatibility, store n_probes to reshape on load - pp_array = np.array(probe_positions).flatten().tolist() - metadata_dict = { - 'General': { - 'title': 'Multislice Wavefunction', - 'signal_type': 'Wavefunction' - }, - 'Simulation': { - 'voltage_eV': float(probe.eV), - 'wavelength_A': float(probe.wavelength), - 'aperture_mrad': float(probe.mrad), - 'probe_positions': pp_array, - 'n_probes': len(probe_positions), + if Dimensions is not None: + dimensions = Dimensions([ + Dimension(name='probe', space='position', + values=np.arange(len(probe_positions))), + Dimension(name='time', space='temporal', units='ps', + values=time_arr), + Dimension(name='kx', space='scattering', units='Å⁻¹', + values=kxs_arr), + Dimension(name='ky', space='scattering', units='Å⁻¹', + values=kys_arr), + Dimension(name='layer', space='position', + values=layer_arr), + ], nav_dimensions=[0, 1], sig_dimensions=[2, 3, 4]) + + # Build metadata from simulation parameters + # Flatten probe_positions for HDF5 compatibility, store n_probes to reshape on load + pp_array = np.array(probe_positions).flatten().tolist() + metadata_dict = { + 'General': { + 'title': 'Multislice Wavefunction', + 'signal_type': 'Wavefunction' + }, + 'Simulation': { + 'voltage_eV': float(probe.eV), + 'wavelength_A': float(probe.wavelength), + 'aperture_mrad': float(probe.mrad), + 'probe_positions': pp_array, + 'n_probes': len(probe_positions), + } } - } - metadata = GeneralMetadata(metadata_dict) - - # Initialize Signal base class (must be called before setting _array - # because Signal.__init__ sets self.data which calls our setter) - super().__init__( - data=None, - name='WFData', - dimensions=dimensions, - signal_type='Diffraction', - metadata=metadata - ) + metadata = Metadata(metadata_dict) # Store array AFTER super().__init__ to avoid being overwritten self._array = array @@ -176,54 +175,6 @@ def __getattr__(self, name): return np.asarray(raw) raise AttributeError(f"'{type(self).__name__}' has no attribute '{name}'") - def to_hdf5_group(self, parent_group, force_datasets=['data'], name=None): - """Override to convert tensors to numpy and exclude non-serializable attrs.""" - # Put numpy data directly in __dict__ for serialization - # (to_dict pulls from __dict__, not properties) - if hasattr(self._array, 'cpu'): - self.__dict__['data'] = self._array.cpu().numpy() - else: - self.__dict__['data'] = np.asarray(self._array) - - # Temporarily remove non-serializable attributes (but keep _array for data property) - orig_probe = self.probe - orig_cache_dir = self.cache_dir - orig_probe_positions = self.probe_positions - orig_kxs = self._kxs - orig_kys = self._kys - orig_time = self._time - orig_xs = self._xs - orig_ys = self._ys - orig_layer = self._layer - - del self.probe - del self.cache_dir - del self.probe_positions - del self._kxs - del self._kys - del self._time - del self._xs - del self._ys - del self._layer - # Note: DON'T delete _array - the data property depends on it for hasattr check - - # Call parent implementation - result = super().to_hdf5_group(parent_group, force_datasets=force_datasets, name=name) - - # Restore all original attributes and clean up - del self.__dict__['data'] # Remove the temp data from __dict__ - self.probe = orig_probe - self.cache_dir = orig_cache_dir - self.probe_positions = orig_probe_positions - self._kxs = orig_kxs - self._kys = orig_kys - self._time = orig_time - self._xs = orig_xs - self._ys = orig_ys - self._layer = orig_layer - - return result - def plot_reciprocal(self,filename=None,whichProbe="mean",whichTimestep="mean",powerscaling=0.25,extent=None,nuke_zerobeam=False): import matplotlib.pyplot as plt fig, ax = plt.subplots() diff --git a/tests/01_potentials_arbitrary.py b/tests/01_potentials_arbitrary.py new file mode 100644 index 0000000..ca087ab --- /dev/null +++ b/tests/01_potentials_arbitrary.py @@ -0,0 +1,33 @@ +import sys,os,time +try: + import pyslice +except ModuleNotFoundError: + print("import failed, falling back to relative paths") + sys.path.insert(0, '../src') +from pyslice.backend import zeros,fft,fftshift +from pyslice import Probe,Propagate,Potential +import matplotlib.pyplot as plt +import numpy as np +try: + import torch + xp = torch +except: + xp = np + +xs=np.linspace(0,11,111) +ys=np.linspace(0,10,101) + +array = zeros((len(xs),len(ys),1))+1+.2*np.sin(10*xs[:,None,None]+13*ys[None,:,None]) +array*=1000 +#plt.imshow(array[:,:,0]) +#plt.show() + +O = Potential(xs, ys, [0], array=array) +P = probe=Probe(xs,ys,mrad=30,eV=100e3) +E = Propagate(P,O) + +#E.plot() +#plt.imshow(np.absolute(E[0,:,:])) +#plt.show() +plt.imshow(np.absolute(fftshift(xp.fft.fft2(E[0,:,:])))) +plt.show() diff --git a/tests/02_unpropagate.py b/tests/02_unpropagate.py new file mode 100644 index 0000000..68b52e8 --- /dev/null +++ b/tests/02_unpropagate.py @@ -0,0 +1,78 @@ +import sys,os +try: + import pyslice +except ModuleNotFoundError: + sys.path.insert(0, '../src') + +from pyslice import Loader,Probe,Propagate,gridFromTrajectory,Potential,differ,to_cpu,calculateObject +import matplotlib.pyplot as plt +import numpy as np + +dump="inputs/hBN_truncated.lammpstrj" +dt=.005 +types={1:"B",2:"N"} + +for flatten in [True,False]: + + # LOAD MD OUTPUT + trajectory=Loader(dump,timestep=dt,atom_mapping=types).load() + xs,ys,zs,lx,ly,lz=gridFromTrajectory(trajectory,sampling=0.1,slice_thickness=0.5) + + # GENERATE PROBE (ENSURE 00_PROBE.PY PASSES BEFORE RUNNING) + xpr=[lx/2,lx/2+5] ; ypr=[ly/2]*2 ; Os=[] + for x,y in zip(xpr,ypr): + + probe=Probe(xs,ys,mrad=5,eV=100e3,probe_xs=[x],probe_ys=[y]) + probe.defocus(200) + #probe.plot() + + # GENERATE THE POTENTIAL (ENSURE 01_POTENTIAL.PY PASSES BEFORE RUNNING) + positions = trajectory.positions[0] + atom_types=trajectory.atom_types + potential = Potential(xs, ys, zs, positions, atom_types, kind="kirkland") + potential.plot() + potential.build() + if flatten: + potential.flatten() # TECHNICALLY calculateObject ONLY RETURNS THE TRULY CORRECT SOLUTION FOR A SINGLE SLICE + p_arry = np.absolute(np.sum(potential.array,axis=2)) + + #fig, ax = plt.subplots() + #ax.imshow(p_arry, cmap="inferno") + #plt.title("Original Potential") + #plt.show() + + # PROPAGATION + # Handle device conversion properly for PyTorch tensors + result = Propagate(probe,potential,onthefly=True) + res = to_cpu(result[0,:,:]) + #fig, ax = plt.subplots() + #ax.imshow(np.absolute(res), cmap="inferno") + #plt.title("|Exit Wave|") + #plt.show() + #fig, ax = plt.subplots() + #ax.imshow(np.angle(res), cmap="inferno") + #plt.title("Phi Exit Wave") + #plt.show() + + # RECALCULATE OBJECT FROM THE RESULT + dO = calculateObject(probe,result[0,:,:],np.zeros((len(xs),len(ys))),weighting=1,dz=0.5) + Os.append(dO) + dO = to_cpu(dO) + fig, ax = plt.subplots() + ax.imshow(np.absolute(dO)**.1, cmap="inferno") + plt.title("Reconstructed Potential") + plt.show() + + O = np.sum(Os,axis=0) + fig, ax = plt.subplots() + ax.imshow(np.absolute(O)**.1, cmap="inferno") + plt.title("summed reconstructed") + plt.show() + + + fig, ax = plt.subplots() + delta = np.absolute(dO-p_arry) + ax.imshow(delta**.1, cmap="inferno") + plt.title("|OP-RP|/|OP|="+str(np.amax(delta)/np.amax(p_arry))) + plt.show() + diff --git a/tests/03_manyprobes.py b/tests/03_manyprobes.py index 5781371..670c594 100644 --- a/tests/03_manyprobes.py +++ b/tests/03_manyprobes.py @@ -21,8 +21,8 @@ # GENERATE PROBE (ENSURE 00_PROBE.PY PASSES BEFORE RUNNING) #probe=Probe(xs,ys,mrad=30,eV=100e3) -xsp = np.linspace(a,3*a,16) -ysp = np.linspace(b,3*b,16) +xsp = np.linspace(10*a-a,10*a-3*a,16) +ysp = np.linspace(10*b-b,10*b-3*b,16) #x,y=np.meshgrid(xs,ys) #xy=np.reshape([x,y],(2,len(x.flat))).T #print(xy) diff --git a/tests/04_haadf.py b/tests/04_haadf.py index d8a4f94..c098e51 100644 --- a/tests/04_haadf.py +++ b/tests/04_haadf.py @@ -29,9 +29,9 @@ # SET UP GRID OF HAADF SCAN POINTS #xy=probe_grid([a,3*a],[b,3*b],14,16) #calculator.setup(trajectory,aperture=30,voltage_eV=100e3,sampling=.1,slice_thickness=.5,probe_positions=xy,cache_levels=[]) -probe_xs = np.linspace(a,3*a,14) -probe_ys = np.linspace(b,3*b,16) -calculator.setup(trajectory,aperture=30,voltage_eV=100e3,sampling=.1,slice_thickness=.5,probe_xs=probe_xs,probe_ys=probe_ys,cache_levels=[]) +probe_xs = np.linspace(10*a-a,10*a-3*a,14) +probe_ys = np.linspace(10*b-b,10*b-3*b,16) +calculator.setup(trajectory,aperture=30,voltage_eV=100e3,sampling=.1,slice_thickness=.5,probe_xs=probe_xs,probe_ys=probe_ys)#,cache_levels=[]) # RUN MULTISLICE exitwaves = calculator.run() @@ -60,5 +60,7 @@ #plt.show() haadf.plot("outputs/figs/04_haadf.png") +haadf.to_sea("04_haadf.sea") + ary=np.asarray(ary) differ(ary[::4,::4],"outputs/haadf-test.npy","HAADF") diff --git a/tests/04_haadf_aberrations.py b/tests/04_haadf_aberrations.py index f09ca2f..b141b53 100644 --- a/tests/04_haadf_aberrations.py +++ b/tests/04_haadf_aberrations.py @@ -29,8 +29,8 @@ # SET UP GRID OF HAADF SCAN POINTS #xy=probe_grid([a,3*a],[b,3*b],14,16) #calculator.setup(trajectory,aperture=30,voltage_eV=100e3,sampling=.1,slice_thickness=.5,probe_positions=xy,cache_levels=[]) -probe_xs = np.linspace(a,3*a,14) -probe_ys = np.linspace(b,3*b,16) +probe_xs = np.linspace(10*a-a,10*a-3*a,14) +probe_ys = np.linspace(10*b-b,10*b-3*b,16) calculator.setup(trajectory,aperture=30,voltage_eV=100e3,sampling=.1,slice_thickness=.5,probe_xs=probe_xs,probe_ys=probe_ys,cache_levels=[]) calculator.base_probe.aberrate({"C30":1e3,"C12":1e2}) # RUN MULTISLICE diff --git a/tests/05_tacaw.py b/tests/05_tacaw.py index 659a0bf..aac3e1c 100644 --- a/tests/05_tacaw.py +++ b/tests/05_tacaw.py @@ -39,6 +39,7 @@ tacaw = TACAWData(exitwaves) +#tacaw.to_sea("05_tacaw.sea") # MANUAL PLOTTING #print(tacaw.frequencies, tacaw.frequencies[65]) #intensity_data = tacaw.intensity[0,65,:,:]**.1 @@ -55,6 +56,8 @@ Z = tacaw.spectral_diffraction(30) ; print(Z.shape) tacaw.plot(Z**.1,"kx","ky",filename="outputs/figs/05_tacaw_30THz.png") +tacaw.to_sea("05_tacaw.sea") + #differ(Z[:,:]**.1,"outputs/tacaw-test.npy","TACAW SLICE") # OR THE DISPERSION: diff --git a/tests/16_sea_eco_integration.py b/tests/16_sea_eco_integration.py index 3922e87..c7e5272 100644 --- a/tests/16_sea_eco_integration.py +++ b/tests/16_sea_eco_integration.py @@ -478,6 +478,143 @@ def test_haadf_is_signal(): return haadf +def test_tacawdata_sea_roundtrip(tacaw): + """Test TACAWData .sea serialization round-trip using PySliceSerial mixin.""" + print("\n--- Testing TACAWData .sea Round-trip (PySliceSerial) ---") + + from pathlib import Path + from pyslice.postprocessing.tacaw_data import TACAWData + from pyslice.backend import to_cpu + + # Store original values + original_array = to_cpu(tacaw._array) + original_freqs = to_cpu(tacaw.frequencies) + original_positions = list(tacaw.probe_positions) + original_cache_dir = tacaw.cache_dir + + # Save using to_sea + output_path = 'outputs/test_tacaw_serial.sea' + tacaw.to_sea(output_path) + assert os.path.exists(output_path) + print(f" Saved TACAWData to: {output_path}") + + # Load using classmethod + loaded = TACAWData.load(output_path) + assert isinstance(loaded, TACAWData) + print(" Loaded TACAWData using TACAWData.load()") + + # Verify data integrity + assert np.allclose(loaded._array, original_array), "Array data mismatch" + print(" PASSED: Array data matches") + + assert np.allclose(loaded.frequencies, original_freqs), "Frequencies mismatch" + print(" PASSED: Frequencies match") + + assert loaded.probe_positions == original_positions, "Probe positions mismatch" + assert isinstance(loaded.probe_positions, list), "Probe positions should be list" + assert isinstance(loaded.probe_positions[0], tuple), "Probe positions elements should be tuples" + print(" PASSED: Probe positions match (list of tuples)") + + assert isinstance(loaded.cache_dir, Path), "cache_dir should be Path" + assert str(loaded.cache_dir) == str(original_cache_dir), "cache_dir mismatch" + print(" PASSED: cache_dir is Path and matches") + + # Verify kxs, kys are restored + assert loaded._kxs is not None, "_kxs should be restored" + assert loaded._kys is not None, "_kys should be restored" + print(" PASSED: Coordinate arrays restored") + + # Verify excluded attrs are None + assert loaded.probe is None, "probe should be None after load" + assert loaded._wf_array is None, "_wf_array should be None after load" + print(" PASSED: Excluded attrs are None") + + print(" PASSED: TACAWData .sea round-trip complete") + return loaded + + +def test_wfdata_sea_roundtrip(wf_data): + """Test WFData .sea serialization round-trip using PySliceSerial mixin.""" + print("\n--- Testing WFData .sea Round-trip (PySliceSerial) ---") + + from pathlib import Path + from pyslice.postprocessing.wf_data import WFData + from pyslice.backend import to_cpu + + # Store original values + original_array = to_cpu(wf_data._array) + original_positions = list(wf_data.probe_positions) + original_kxs = to_cpu(wf_data.kxs) + + # Save using to_sea + output_path = 'outputs/test_wf_serial.sea' + wf_data.to_sea(output_path) + assert os.path.exists(output_path) + print(f" Saved WFData to: {output_path}") + + # Load using classmethod + loaded = WFData.load(output_path) + assert isinstance(loaded, WFData) + print(" Loaded WFData using WFData.load()") + + # Verify data integrity + assert np.allclose(loaded._array, original_array), "Array data mismatch" + print(" PASSED: Array data matches") + + assert loaded.probe_positions == original_positions, "Probe positions mismatch" + assert isinstance(loaded.probe_positions, list), "Probe positions should be list" + print(" PASSED: Probe positions match") + + assert np.allclose(loaded.kxs, original_kxs), "kxs mismatch" + print(" PASSED: Coordinate arrays match") + + assert isinstance(loaded.cache_dir, Path), "cache_dir should be Path" + print(" PASSED: cache_dir is Path") + + print(" PASSED: WFData .sea round-trip complete") + return loaded + + +def test_haadfdata_sea_roundtrip(haadf): + """Test HAADFData .sea serialization round-trip using PySliceSerial mixin.""" + print("\n--- Testing HAADFData .sea Round-trip (PySliceSerial) ---") + + from pathlib import Path + from pyslice.postprocessing.haadf_data import HAADFData + from pyslice.backend import to_cpu + + # Store original values + original_array = to_cpu(haadf._array) + original_positions = list(haadf.probe_positions) + + # Save using to_sea + output_path = 'outputs/test_haadf_serial.sea' + haadf.to_sea(output_path) + assert os.path.exists(output_path) + print(f" Saved HAADFData to: {output_path}") + + # Load using classmethod + loaded = HAADFData.load(output_path) + assert isinstance(loaded, HAADFData) + print(" Loaded HAADFData using HAADFData.load()") + + # Verify data integrity + assert np.allclose(loaded._array, original_array), "Array data mismatch" + print(" PASSED: Array data matches") + + # HAADFData stores probe_positions as array, compare as arrays + loaded_positions = np.array(loaded.probe_positions) if isinstance(loaded.probe_positions, list) else loaded.probe_positions + original_positions_arr = np.array(original_positions) + assert np.allclose(loaded_positions, original_positions_arr), "Probe positions mismatch" + print(" PASSED: Probe positions match") + + assert isinstance(loaded.cache_dir, Path), "cache_dir should be Path" + print(" PASSED: cache_dir is Path") + + print(" PASSED: HAADFData .sea round-trip complete") + return loaded + + if __name__ == '__main__': # Run all tests print("=" * 60) @@ -499,6 +636,13 @@ def test_haadf_is_signal(): tacaw = test_tacawdata_is_signal(wf_data) haadf = test_haadf_is_signal() + print("\n" + "=" * 60) + print("PART 3: PySliceSerial Mixin Tests (.sea round-trips)") + print("=" * 60) + test_tacawdata_sea_roundtrip(tacaw) + test_wfdata_sea_roundtrip(wf_data) + test_haadfdata_sea_roundtrip(haadf) + print("\n" + "=" * 60) print("All sea-eco integration tests PASSED!") print("=" * 60) diff --git a/tests/16b_sea_eco_reload.py b/tests/16b_sea_eco_reload.py new file mode 100644 index 0000000..1ac10a1 --- /dev/null +++ b/tests/16b_sea_eco_reload.py @@ -0,0 +1,40 @@ +import os,sys,shutil +import matplotlib.pyplot as plt + +# ensure we have pySEA downloaded +if not os.path.exists("../../pySEA"): + print("downloading pySEA from private github repo") + u,k=open("/home/qwe/.gitp").readlines() # may need git username and git key for private repo + u=u.strip() ; k=k.strip() + os.system("git clone https://"+u+":"+k+"@github.com/sea-ecosystem/sea-eco") + shutil.move("sea-eco/src/pySEA","../../") + shutil.rmtree("sea-eco") + +# make sure the import works before we try anything else.... +sys.path.insert(1,"../../") +from pySEA.sea_eco.io import load +from pySEA.sea_eco.architecture.base_structure_numpy import SEAFile + +# run 05, which should conditionally save a .sea file if pySEA is found +if not os.path.exists("04_haadf.sea"): + print("sea file does not exist, running 04_haadf") + os.system("python3 04_haadf.py") + +# run 05, which should conditionally save a .sea file if pySEA is found +if not os.path.exists("05_tacaw.sea"): + print("sea file does not exist, running 05_tacaw") + os.system("python3 05_tacaw.py") + +# attempt to reload and plot it +loaded = load(file_path='05_tacaw.sea') +loaded.show_tree() + +loaded.show(dims=('frequency','kx')) +plt.show() + +loaded.show(dims=('kx','ky')) +plt.show() + +loaded = load(file_path='04_haadf.sea') +loaded.show() +plt.show() diff --git a/tests/17_FPcifCBED.py b/tests/17_FPcifCBED.py index e9eb34a..f289cb0 100644 --- a/tests/17_FPcifCBED.py +++ b/tests/17_FPcifCBED.py @@ -4,12 +4,12 @@ except ModuleNotFoundError: sys.path.insert(0, '../src') -from pyslice import Loader,gridFromTrajectory,Potential,MultisliceCalculator,HAADFData,differ +from pyslice import Loader,gridFromTrajectory,Potential,MultisliceCalculator,HAADFData import matplotlib.pyplot as plt import numpy as np -cif = "inputs/hBN.cif" ; a,b=2.4907733333333337,2.1570729817355123 +cif = "inputs/hBN_cif.cif" ; a,b=2.4907733333333337,2.1570729817355123 trajectory=Loader(cif).load() trajectory = trajectory.tile_positions([5,5,1]) trajectory = trajectory.generate_random_displacements(n_displacements=20,sigma=.1,seed=0) # PRO TIP: use a random seed so you get repeatable "random" configurations and don't need to re-multislice on subsequent runs @@ -38,4 +38,4 @@ array = exitwaves.array -print("array shape",array.shape) \ No newline at end of file +print("array shape",array.shape) diff --git a/tests/18_caching.py b/tests/18_caching.py index 9d0a0a5..b79aa7c 100644 --- a/tests/18_caching.py +++ b/tests/18_caching.py @@ -4,7 +4,7 @@ except ModuleNotFoundError: sys.path.insert(0, '../src') -from pyslice import Loader,probe_grid,MultisliceCalculator,HAADFData,differ +from pyslice import Loader,MultisliceCalculator,HAADFData,differ import numpy as np import matplotlib.pyplot as plt @@ -51,8 +51,8 @@ print("3. one timestep, many probes, normal caching") traj3=trajectory.get_random_timesteps(1,seed=3) calculator=MultisliceCalculator() - probe_xs = np.linspace(a,3*a,14) - probe_ys = np.linspace(b,3*b,16) + probe_xs = np.linspace(10*a-a,10*a-3*a,14) + probe_ys = np.linspace(10*b-b,10*b-3*b,16) calculator.setup(traj3,aperture=30,voltage_eV=100e3,sampling=.1,slice_thickness=.5,probe_xs=probe_xs,probe_ys=probe_ys) exitwaves = calculator.run() differ(exitwaves.array[::5,:,::5,::5,:],"outputs/caching/03-test.npy","03") # p,t,x,y,l indices @@ -71,8 +71,8 @@ print("5. many timesteps, many probes, normal caching") traj5=trajectory.get_random_timesteps(5,seed=5) calculator=MultisliceCalculator() - probe_xs = np.linspace(a,3*a,9) - probe_ys = np.linspace(b,3*b,10) + probe_xs = np.linspace(10*a-a,10*a-3*a,9) + probe_ys = np.linspace(10*b-b,10*b-3*b,10) calculator.setup(traj5,aperture=30,voltage_eV=100e3,sampling=.1,slice_thickness=.5,probe_xs=probe_xs,probe_ys=probe_ys) exitwaves = calculator.run() differ(exitwaves.array[:,:,::10,::10,:],"outputs/caching/05-test.npy","05") # p,t,x,y,l indices @@ -82,8 +82,8 @@ print("5.5 many timesteps, many probes, normal caching, with decoherence added") traj55=trajectory.get_random_timesteps(5,seed=5) calculator=MultisliceCalculator() - probe_xs = np.linspace(a,3*a,9) - probe_ys = np.linspace(b,3*b,10) + probe_xs = np.linspace(10*a-a,10*a-3*a,9) + probe_ys = np.linspace(10*b-b,10*b-3*b,10) calculator.setup(traj55,aperture=30,voltage_eV=100e3,sampling=.1,slice_thickness=.5,probe_xs=probe_xs,probe_ys=probe_ys) calculator.base_probe.addSpatialDecoherence(1000,7) exitwaves = calculator.run() @@ -95,8 +95,8 @@ print("6. many timesteps, many probes, no caching") traj6=trajectory.get_random_timesteps(5,seed=6) calculator=MultisliceCalculator() - probe_xs = np.linspace(a,3*a,9) - probe_ys = np.linspace(b,3*b,10) + probe_xs = np.linspace(10*a-a,10*a-3*a,9) + probe_ys = np.linspace(10*b-b,10*b-3*b,10) calculator.setup(traj6,aperture=30,voltage_eV=100e3,sampling=.1,slice_thickness=.5,probe_xs=probe_xs,probe_ys=probe_ys,cache_levels=[]) exitwaves = calculator.run() differ(exitwaves.array[:,:,::10,::10,:],"outputs/caching/06-test.npy","06") # p,t,x,y,l indices @@ -115,8 +115,8 @@ print("8. many timesteps, many probes, layerwise caching") traj8=trajectory.get_random_timesteps(5,seed=8) calculator=MultisliceCalculator() - probe_xs = np.linspace(a,3*a,6) - probe_ys = np.linspace(b,3*b,7) + probe_xs = np.linspace(10*a-a,10*a-3*a,6) + probe_ys = np.linspace(10*b-b,10*b-3*b,7) calculator.setup(traj8,aperture=30,voltage_eV=100e3,sampling=.1,slice_thickness=.5,probe_xs=probe_xs,probe_ys=probe_ys,cache_levels=["slices"]) exitwaves = calculator.run() differ(exitwaves.array[:,::3,::20,::20,::5],"outputs/caching/08-test.npy","08") # p,t,x,y,l indices @@ -135,8 +135,8 @@ print("10. one timestep, many probes, layerwise caching") traj10=trajectory.get_random_timesteps(1,seed=10) calculator=MultisliceCalculator() - probe_xs = np.linspace(a,3*a,9) - probe_ys = np.linspace(b,3*b,10) + probe_xs = np.linspace(10*a-a,10*a-3*a,9) + probe_ys = np.linspace(10*b-b,10*b-3*b,10) calculator.setup(traj10,aperture=30,voltage_eV=100e3,sampling=.1,slice_thickness=.5,probe_xs=probe_xs,probe_ys=probe_ys,cache_levels=["slices"]) exitwaves = calculator.run() differ(exitwaves.array[:,:,::10,::10,::5],"outputs/caching/10-test.npy","10") # p,t,x,y,l indices diff --git a/tests/20_probepositions.py b/tests/20_probepositions.py new file mode 100644 index 0000000..c6ad4fd --- /dev/null +++ b/tests/20_probepositions.py @@ -0,0 +1,61 @@ +import sys,os +try: + import pyslice +except ModuleNotFoundError: + sys.path.insert(0, '../src') + +from pyslice import Loader,MultisliceCalculator,HAADFData,differ,gridFromTrajectory,Potential + +import numpy as np +import matplotlib.pyplot as plt +import os,shutil + +#if os.path.exists("psi_data"): +# shutil.rmtree("psi_data") + +dump="inputs/hBN_truncated.lammpstrj" +dt=.005 +types={1:"B",2:"N"} +a,b=2.4907733333333337,2.1570729817355123 + +# LOAD TRAJECTORY +trajectory=Loader(dump,timestep=dt,atom_mapping=types).load() +# SELECT 10 "RANDOM" TIMESTEPS (use seed for reproducibility) +trajectory=trajectory.get_random_timesteps(3,seed=5) +# REPLACE A SINGLE ATOM WITH SILICON +dxyz = trajectory.positions[0,:,:]-np.asarray([a,b*4/3,0])[None,:] +distances = np.sqrt(np.sum((dxyz)**2,axis=1)) +i = np.argmin(distances) # which atom is closest to a,b? +trajectory.atom_types[i] = 28 +# PREVIEW IT: +xs,ys,zs,lx,ly,lz=gridFromTrajectory(trajectory,sampling=0.1,slice_thickness=0.5) +potential = Potential(xs, ys, zs, trajectory.positions[0], trajectory.atom_types, kind="kirkland") +potential.plot() +# TRIM TO 10x10 UC +trajectory=trajectory.slice_positions([0,10*a],[0,10*b]) +# CHECK AGAIN, POST-CROP +xs,ys,zs,lx,ly,lz=gridFromTrajectory(trajectory,sampling=0.1,slice_thickness=0.5) +potential = Potential(xs, ys, zs, trajectory.positions[0], trajectory.atom_types, kind="kirkland") +potential.plot() +# CREATE CALCULATOR OBJECT +calculator=MultisliceCalculator() +# SET UP GRID OF HAADF SCAN POINTS +#xy=probe_grid([a,3*a],[b,3*b],14,16) +#calculator.setup(trajectory,aperture=30,voltage_eV=100e3,sampling=.1,slice_thickness=.5,probe_positions=xy,cache_levels=[]) +probe_xs = np.linspace(a/2,2.5*a,14) +probe_ys = np.linspace(b/2,2.5*b,16) +calculator.setup(trajectory,aperture=30,voltage_eV=100e3,sampling=.1,slice_thickness=.5,probe_xs=probe_xs,probe_ys=probe_ys,cache_levels=[]) +# RUN MULTISLICE +exitwaves = calculator.run() + +haadf=HAADFData(exitwaves) +ary=haadf.calculateADF(preview=False) # use preview=True to view the collection angles of the ADF detector in reciprocal space +xs=haadf.xs ; ys=haadf.ys + +#fig, ax = plt.subplots() +#ax.imshow(ary.T, cmap="inferno") +#plt.show() +haadf.plot()#"outputs/figs/04_haadf.png") + +#ary=np.asarray(ary) +#differ(ary[::4,::4],"outputs/haadf-test.npy","HAADF") diff --git a/tests/cpu.sberr b/tests/cpu.sberr index 3c45608..ad601e6 100644 --- a/tests/cpu.sberr +++ b/tests/cpu.sberr @@ -6,4 +6,5 @@ cuda/11.4.0: Please note, CUDA/11.4.0 is required if compiling for NVIDIA K80 GPUs. Newer toolchain versions will not work for K80's. --------------------------------------------------------------------------------DeprecationWarning: 'source deactivate' is deprecated. Use 'conda deactivate'. +-------------------------------------------------------------------------------runAllTests.sh: line 5: 2727913 Killed python3 $s >> runAllTests-$1.log 2>> runAllTests-$1.log +slurmstepd: error: Detected 1 oom_kill event in StepId=4581718.batch. Some of the step tasks have been OOM Killed. diff --git a/tests/gpu.sberr b/tests/gpu.sberr index 3c45608..dd3d051 100644 --- a/tests/gpu.sberr +++ b/tests/gpu.sberr @@ -6,4 +6,4 @@ cuda/11.4.0: Please note, CUDA/11.4.0 is required if compiling for NVIDIA K80 GPUs. Newer toolchain versions will not work for K80's. --------------------------------------------------------------------------------DeprecationWarning: 'source deactivate' is deprecated. Use 'conda deactivate'. +------------------------------------------------------------------------------- \ No newline at end of file diff --git a/tests/numpy.sberr b/tests/numpy.sberr index 3c45608..dd3d051 100644 --- a/tests/numpy.sberr +++ b/tests/numpy.sberr @@ -6,4 +6,4 @@ cuda/11.4.0: Please note, CUDA/11.4.0 is required if compiling for NVIDIA K80 GPUs. Newer toolchain versions will not work for K80's. --------------------------------------------------------------------------------DeprecationWarning: 'source deactivate' is deprecated. Use 'conda deactivate'. +------------------------------------------------------------------------------- \ No newline at end of file diff --git a/tests/outputs/caching/01-test.npy b/tests/outputs/caching/01-test.npy index ce47ef2..9490399 100644 Binary files a/tests/outputs/caching/01-test.npy and b/tests/outputs/caching/01-test.npy differ diff --git a/tests/outputs/caching/02-test.npy b/tests/outputs/caching/02-test.npy index 800602f..1886695 100644 Binary files a/tests/outputs/caching/02-test.npy and b/tests/outputs/caching/02-test.npy differ diff --git a/tests/outputs/caching/03-test.npy b/tests/outputs/caching/03-test.npy index 20d5ae9..2670fbe 100644 Binary files a/tests/outputs/caching/03-test.npy and b/tests/outputs/caching/03-test.npy differ diff --git a/tests/outputs/caching/04-test.npy b/tests/outputs/caching/04-test.npy index 6671bad..e073e83 100644 Binary files a/tests/outputs/caching/04-test.npy and b/tests/outputs/caching/04-test.npy differ diff --git a/tests/outputs/caching/05-test.npy b/tests/outputs/caching/05-test.npy index a587d58..7d63d00 100644 Binary files a/tests/outputs/caching/05-test.npy and b/tests/outputs/caching/05-test.npy differ diff --git a/tests/outputs/caching/06-test.npy b/tests/outputs/caching/06-test.npy index bef999a..ee51f5b 100644 Binary files a/tests/outputs/caching/06-test.npy and b/tests/outputs/caching/06-test.npy differ diff --git a/tests/outputs/caching/07-test.npy b/tests/outputs/caching/07-test.npy index 42cf408..c49aca6 100644 Binary files a/tests/outputs/caching/07-test.npy and b/tests/outputs/caching/07-test.npy differ diff --git a/tests/outputs/caching/08-test.npy b/tests/outputs/caching/08-test.npy index d589472..fa04230 100644 Binary files a/tests/outputs/caching/08-test.npy and b/tests/outputs/caching/08-test.npy differ diff --git a/tests/outputs/caching/09-test.npy b/tests/outputs/caching/09-test.npy index 34eec94..28a73dd 100644 Binary files a/tests/outputs/caching/09-test.npy and b/tests/outputs/caching/09-test.npy differ diff --git a/tests/outputs/caching/10-test.npy b/tests/outputs/caching/10-test.npy index aad9c97..7e9e096 100644 Binary files a/tests/outputs/caching/10-test.npy and b/tests/outputs/caching/10-test.npy differ diff --git a/tests/outputs/defocus-test.npy b/tests/outputs/defocus-test.npy index a1e141e..4315045 100644 Binary files a/tests/outputs/defocus-test.npy and b/tests/outputs/defocus-test.npy differ diff --git a/tests/outputs/figs/00_probe.png b/tests/outputs/figs/00_probe.png index 7f9eff0..b01f38b 100644 Binary files a/tests/outputs/figs/00_probe.png and b/tests/outputs/figs/00_probe.png differ diff --git a/tests/outputs/figs/00_probe_0.png b/tests/outputs/figs/00_probe_0.png index 5ce47dd..7d388c7 100644 Binary files a/tests/outputs/figs/00_probe_0.png and b/tests/outputs/figs/00_probe_0.png differ diff --git a/tests/outputs/figs/00_probe_1.png b/tests/outputs/figs/00_probe_1.png index df924ee..532cdd6 100644 Binary files a/tests/outputs/figs/00_probe_1.png and b/tests/outputs/figs/00_probe_1.png differ diff --git a/tests/outputs/figs/00_probe_2.png b/tests/outputs/figs/00_probe_2.png index bcf56ae..b19a613 100644 Binary files a/tests/outputs/figs/00_probe_2.png and b/tests/outputs/figs/00_probe_2.png differ diff --git a/tests/outputs/figs/00_probe_3.png b/tests/outputs/figs/00_probe_3.png index 8079b3a..dd476af 100644 Binary files a/tests/outputs/figs/00_probe_3.png and b/tests/outputs/figs/00_probe_3.png differ diff --git a/tests/outputs/figs/00_probe_4.png b/tests/outputs/figs/00_probe_4.png index bf82b94..71aeb6e 100644 Binary files a/tests/outputs/figs/00_probe_4.png and b/tests/outputs/figs/00_probe_4.png differ diff --git a/tests/outputs/figs/01_potentials.png b/tests/outputs/figs/01_potentials.png index 0d3da73..b349266 100644 Binary files a/tests/outputs/figs/01_potentials.png and b/tests/outputs/figs/01_potentials.png differ diff --git a/tests/outputs/figs/02_propagate_otf=False.png b/tests/outputs/figs/02_propagate_otf=False.png index 10d707a..4c3a0dd 100644 Binary files a/tests/outputs/figs/02_propagate_otf=False.png and b/tests/outputs/figs/02_propagate_otf=False.png differ diff --git a/tests/outputs/figs/02_propagate_otf=True.png b/tests/outputs/figs/02_propagate_otf=True.png index 10d707a..4c3a0dd 100644 Binary files a/tests/outputs/figs/02_propagate_otf=True.png and b/tests/outputs/figs/02_propagate_otf=True.png differ diff --git a/tests/outputs/figs/03_manyprobes.png b/tests/outputs/figs/03_manyprobes.png index ac09183..5898736 100644 Binary files a/tests/outputs/figs/03_manyprobes.png and b/tests/outputs/figs/03_manyprobes.png differ diff --git a/tests/outputs/figs/04_haadf.png b/tests/outputs/figs/04_haadf.png index 7ad5bf0..ce8595f 100644 Binary files a/tests/outputs/figs/04_haadf.png and b/tests/outputs/figs/04_haadf.png differ diff --git a/tests/outputs/figs/04_haadf_aberrations.png b/tests/outputs/figs/04_haadf_aberrations.png index 25eb8eb..a3df7b8 100644 Binary files a/tests/outputs/figs/04_haadf_aberrations.png and b/tests/outputs/figs/04_haadf_aberrations.png differ diff --git a/tests/outputs/figs/04_haadf_cbed.png b/tests/outputs/figs/04_haadf_cbed.png index 32229e2..c1aca5a 100644 Binary files a/tests/outputs/figs/04_haadf_cbed.png and b/tests/outputs/figs/04_haadf_cbed.png differ diff --git a/tests/outputs/figs/05_tacaw_30THz.png b/tests/outputs/figs/05_tacaw_30THz.png index 098d01a..2ca767c 100644 Binary files a/tests/outputs/figs/05_tacaw_30THz.png and b/tests/outputs/figs/05_tacaw_30THz.png differ diff --git a/tests/outputs/figs/05_tacaw_diff.png b/tests/outputs/figs/05_tacaw_diff.png index 642f888..657a6ef 100644 Binary files a/tests/outputs/figs/05_tacaw_diff.png and b/tests/outputs/figs/05_tacaw_diff.png differ diff --git a/tests/outputs/figs/05_tacaw_disp.png b/tests/outputs/figs/05_tacaw_disp.png index 59d8b4a..76b604f 100644 Binary files a/tests/outputs/figs/05_tacaw_disp.png and b/tests/outputs/figs/05_tacaw_disp.png differ diff --git a/tests/outputs/figs/05_tacawcrop_30THz.png b/tests/outputs/figs/05_tacawcrop_30THz.png index 43d1c9d..71a62c2 100644 Binary files a/tests/outputs/figs/05_tacawcrop_30THz.png and b/tests/outputs/figs/05_tacawcrop_30THz.png differ diff --git a/tests/outputs/figs/05_tacawcrop_diff.png b/tests/outputs/figs/05_tacawcrop_diff.png index c73925a..349eefb 100644 Binary files a/tests/outputs/figs/05_tacawcrop_diff.png and b/tests/outputs/figs/05_tacawcrop_diff.png differ diff --git a/tests/outputs/figs/05_tacawcrop_disp.png b/tests/outputs/figs/05_tacawcrop_disp.png index 59d8b4a..f0d8af0 100644 Binary files a/tests/outputs/figs/05_tacawcrop_disp.png and b/tests/outputs/figs/05_tacawcrop_disp.png differ diff --git a/tests/outputs/figs/06_loaders_0.png b/tests/outputs/figs/06_loaders_0.png index 132b352..7451d98 100644 Binary files a/tests/outputs/figs/06_loaders_0.png and b/tests/outputs/figs/06_loaders_0.png differ diff --git a/tests/outputs/figs/06_loaders_1.png b/tests/outputs/figs/06_loaders_1.png index eedf689..18d960a 100644 Binary files a/tests/outputs/figs/06_loaders_1.png and b/tests/outputs/figs/06_loaders_1.png differ diff --git a/tests/outputs/figs/06_loaders_2.png b/tests/outputs/figs/06_loaders_2.png index 0d3da73..b349266 100644 Binary files a/tests/outputs/figs/06_loaders_2.png and b/tests/outputs/figs/06_loaders_2.png differ diff --git a/tests/outputs/figs/06_loaders_3.png b/tests/outputs/figs/06_loaders_3.png index 88ca4cf..a412e66 100644 Binary files a/tests/outputs/figs/06_loaders_3.png and b/tests/outputs/figs/06_loaders_3.png differ diff --git a/tests/outputs/figs/07_defocus_2D.png b/tests/outputs/figs/07_defocus_2D.png index da1cc8f..c752ec8 100644 Binary files a/tests/outputs/figs/07_defocus_2D.png and b/tests/outputs/figs/07_defocus_2D.png differ diff --git a/tests/outputs/figs/07_defocus_3D.png b/tests/outputs/figs/07_defocus_3D.png index f32c54d..4b74012 100644 Binary files a/tests/outputs/figs/07_defocus_3D.png and b/tests/outputs/figs/07_defocus_3D.png differ diff --git a/tests/outputs/figs/08_LACBED_iterative.png b/tests/outputs/figs/08_LACBED_iterative.png index 75362b8..2fde38f 100644 Binary files a/tests/outputs/figs/08_LACBED_iterative.png and b/tests/outputs/figs/08_LACBED_iterative.png differ diff --git a/tests/outputs/figs/08_LACBED_onthefly.png b/tests/outputs/figs/08_LACBED_onthefly.png index 4873e22..70fab28 100644 Binary files a/tests/outputs/figs/08_LACBED_onthefly.png and b/tests/outputs/figs/08_LACBED_onthefly.png differ diff --git a/tests/outputs/figs/10_midgley_30THz.png b/tests/outputs/figs/10_midgley_30THz.png index 9e22cdf..99726a9 100644 Binary files a/tests/outputs/figs/10_midgley_30THz.png and b/tests/outputs/figs/10_midgley_30THz.png differ diff --git a/tests/outputs/figs/10_midgley_diff.png b/tests/outputs/figs/10_midgley_diff.png index b374bb9..e775265 100644 Binary files a/tests/outputs/figs/10_midgley_diff.png and b/tests/outputs/figs/10_midgley_diff.png differ diff --git a/tests/outputs/figs/10_midgley_diff2.png b/tests/outputs/figs/10_midgley_diff2.png index 56b6e59..825d154 100644 Binary files a/tests/outputs/figs/10_midgley_diff2.png and b/tests/outputs/figs/10_midgley_diff2.png differ diff --git a/tests/outputs/figs/10_midgley_disp.png b/tests/outputs/figs/10_midgley_disp.png index 3a9cc08..3312bdd 100644 Binary files a/tests/outputs/figs/10_midgley_disp.png and b/tests/outputs/figs/10_midgley_disp.png differ diff --git a/tests/outputs/figs/11_SED_30THz.png b/tests/outputs/figs/11_SED_30THz.png index 1245159..f2c9c73 100644 Binary files a/tests/outputs/figs/11_SED_30THz.png and b/tests/outputs/figs/11_SED_30THz.png differ diff --git a/tests/outputs/figs/12_aberrations_espread.png b/tests/outputs/figs/12_aberrations_espread.png index 3dcc7f9..166ea3d 100644 Binary files a/tests/outputs/figs/12_aberrations_espread.png and b/tests/outputs/figs/12_aberrations_espread.png differ diff --git a/tests/outputs/figs/12_aberrations_tableau.png b/tests/outputs/figs/12_aberrations_tableau.png index a6b4735..ca63f48 100644 Binary files a/tests/outputs/figs/12_aberrations_tableau.png and b/tests/outputs/figs/12_aberrations_tableau.png differ diff --git a/tests/outputs/figs/17_cbed.png b/tests/outputs/figs/17_cbed.png index 5bf350f..32a9e59 100644 Binary files a/tests/outputs/figs/17_cbed.png and b/tests/outputs/figs/17_cbed.png differ diff --git a/tests/outputs/figs/17_haadf.png b/tests/outputs/figs/17_haadf.png index bc9c713..05936a0 100644 Binary files a/tests/outputs/figs/17_haadf.png and b/tests/outputs/figs/17_haadf.png differ diff --git a/tests/outputs/figs/17_potential.png b/tests/outputs/figs/17_potential.png index f4f1527..0cb81f2 100644 Binary files a/tests/outputs/figs/17_potential.png and b/tests/outputs/figs/17_potential.png differ diff --git a/tests/outputs/figs/19_coherence_dZ1.png b/tests/outputs/figs/19_coherence_dZ1.png new file mode 100644 index 0000000..001ea15 Binary files /dev/null and b/tests/outputs/figs/19_coherence_dZ1.png differ diff --git a/tests/outputs/figs/19_coherence_dZ2.png b/tests/outputs/figs/19_coherence_dZ2.png new file mode 100644 index 0000000..48c2a34 Binary files /dev/null and b/tests/outputs/figs/19_coherence_dZ2.png differ diff --git a/tests/outputs/figs/19_coherence_eV1.png b/tests/outputs/figs/19_coherence_eV1.png new file mode 100644 index 0000000..d6727bb Binary files /dev/null and b/tests/outputs/figs/19_coherence_eV1.png differ diff --git a/tests/outputs/figs/19_coherence_eV2.png b/tests/outputs/figs/19_coherence_eV2.png new file mode 100644 index 0000000..1ee0bee Binary files /dev/null and b/tests/outputs/figs/19_coherence_eV2.png differ diff --git a/tests/outputs/figs/19_coherence_eVdZ.png b/tests/outputs/figs/19_coherence_eVdZ.png new file mode 100644 index 0000000..bc18873 Binary files /dev/null and b/tests/outputs/figs/19_coherence_eVdZ.png differ diff --git a/tests/outputs/figs/19_coherence_vanilla.png b/tests/outputs/figs/19_coherence_vanilla.png new file mode 100644 index 0000000..9561384 Binary files /dev/null and b/tests/outputs/figs/19_coherence_vanilla.png differ diff --git a/tests/outputs/haadf-test.npy b/tests/outputs/haadf-test.npy index 799df5b..f9fa99e 100644 Binary files a/tests/outputs/haadf-test.npy and b/tests/outputs/haadf-test.npy differ diff --git a/tests/outputs/haadfaberrations-test.npy b/tests/outputs/haadfaberrations-test.npy index 9ef1c16..9a4a79d 100644 Binary files a/tests/outputs/haadfaberrations-test.npy and b/tests/outputs/haadfaberrations-test.npy differ diff --git a/tests/outputs/lacbed-test.npy b/tests/outputs/lacbed-test.npy index 9cb51b9..f54aa5f 100644 Binary files a/tests/outputs/lacbed-test.npy and b/tests/outputs/lacbed-test.npy differ diff --git a/tests/outputs/manyprobes-test.npy b/tests/outputs/manyprobes-test.npy index 74c4114..b7296a5 100644 Binary files a/tests/outputs/manyprobes-test.npy and b/tests/outputs/manyprobes-test.npy differ diff --git a/tests/outputs/potentials-test.npy b/tests/outputs/potentials-test.npy index f22f932..7aac50e 100644 Binary files a/tests/outputs/potentials-test.npy and b/tests/outputs/potentials-test.npy differ diff --git a/tests/outputs/probe-test.npy b/tests/outputs/probe-test.npy index 522d35d..3e66e1d 100644 Binary files a/tests/outputs/probe-test.npy and b/tests/outputs/probe-test.npy differ diff --git a/tests/outputs/propagate-test.npy b/tests/outputs/propagate-test.npy index 91d2965..8836b9d 100644 Binary files a/tests/outputs/propagate-test.npy and b/tests/outputs/propagate-test.npy differ diff --git a/tests/outputs/tacaw-test.npy b/tests/outputs/tacaw-test.npy index 2e37a1f..b660d0b 100644 Binary files a/tests/outputs/tacaw-test.npy and b/tests/outputs/tacaw-test.npy differ diff --git a/tests/outputs/tacawcrop-test.npy b/tests/outputs/tacawcrop-test.npy index 2d72f9b..56780d5 100644 Binary files a/tests/outputs/tacawcrop-test.npy and b/tests/outputs/tacawcrop-test.npy differ diff --git a/tests/outputs/test_complex_signal.sea b/tests/outputs/test_complex_signal.sea new file mode 100644 index 0000000..5959c8d Binary files /dev/null and b/tests/outputs/test_complex_signal.sea differ diff --git a/tests/outputs/test_sea_eco.sea b/tests/outputs/test_sea_eco.sea new file mode 100644 index 0000000..ecdb220 Binary files /dev/null and b/tests/outputs/test_sea_eco.sea differ diff --git a/tests/outputs/test_signal_set.sea b/tests/outputs/test_signal_set.sea new file mode 100644 index 0000000..ec3e5c7 Binary files /dev/null and b/tests/outputs/test_signal_set.sea differ diff --git a/tests/runAllTests-numpy.log b/tests/runAllTests-numpy.log index 66c7ec1..b2ce330 100644 --- a/tests/runAllTests-numpy.log +++ b/tests/runAllTests-numpy.log @@ -1,16 +1,18 @@ -Sat Jan 10 17:08:00 EST 2026 +Wed Feb 4 12:50:08 EST 2026 python3 00_probe.py import failed, falling back to relative paths PyTorch not available, falling back to NumPy PyTorch not available, falling back to NumPy +failed to import pySEA: No module named 'pySEA' PyTorch not available, falling back to NumPy PyTorch not available, falling back to NumPy PyTorch not available, falling back to NumPy -probe import took 2.1932120323181152 s -differ import took 0.0002701282501220703 s +probe import took 1.6938552856445312 s +differ import took 0.00021028518676757812 s python3 01_potentials.py PyTorch not available, falling back to NumPy PyTorch not available, falling back to NumPy +failed to import pySEA: No module named 'pySEA' PyTorch not available, falling back to NumPy PyTorch not available, falling back to NumPy PyTorch not available, falling back to NumPy @@ -18,18 +20,21 @@ PyTorch not available, falling back to NumPy python3 02_propagate_otf=False.py PyTorch not available, falling back to NumPy PyTorch not available, falling back to NumPy +failed to import pySEA: No module named 'pySEA' PyTorch not available, falling back to NumPy PyTorch not available, falling back to NumPy PyTorch not available, falling back to NumPy python3 02_propagate_otf=True.py PyTorch not available, falling back to NumPy PyTorch not available, falling back to NumPy +failed to import pySEA: No module named 'pySEA' PyTorch not available, falling back to NumPy PyTorch not available, falling back to NumPy PyTorch not available, falling back to NumPy python3 03_manyprobes.py PyTorch not available, falling back to NumPy PyTorch not available, falling back to NumPy +failed to import pySEA: No module named 'pySEA' PyTorch not available, falling back to NumPy PyTorch not available, falling back to NumPy PyTorch not available, falling back to NumPy @@ -37,61 +42,65 @@ PyTorch not available, falling back to NumPy python3 04_haadf.py PyTorch not available, falling back to NumPy PyTorch not available, falling back to NumPy +failed to import pySEA: No module named 'pySEA' PyTorch not available, falling back to NumPy PyTorch not available, falling back to NumPy PyTorch not available, falling back to NumPy - Processing frames: 0%| | 0/3 [00:00 wfdata (1, 100, 1495, 1295, 1) + Processing frames: 0%| | 0/100 [00:00 wfdata (1, 100, 1495, 1295, 1) + Processing frames: 0%| | 0/100 [00:00 wfdata (1, 100, 481, 481, 1) + Processing frames: 0%| | 0/100 [00:00 wfdata (1, 100, 1495, 1295, 1) + Processing frames: 0%| | 0/100 [00:00 wfdata torch.Size([1, 100, 1495, 1295, 1]) +failed to import pySEA: No module named 'pySEA' + Processing frames: 0%| | 0/100 [00:00 wfdata torch.Size([1, 100, 1495, 1295, 1]) +failed to import pySEA: No module named 'pySEA' + Processing frames: 0%| | 0/100 [00:00 wfdata torch.Size([1, 100, 481, 481, 1]) +failed to import pySEA: No module named 'pySEA' + Processing frames: 0%| | 0/100 [00:00 wfdata torch.Size([1, 100, 1495, 1295, 1]) +failed to import pySEA: No module named 'pySEA' + Processing frames: 0%| | 0/100 [00:00 wfdata torch.Size([1, 100, 1495, 1295, 1]) +failed to import pySEA: No module named 'pySEA' + Processing frames: 0%| | 0/100 [00:00 wfdata torch.Size([1, 100, 1495, 1295, 1]) +failed to import pySEA: No module named 'pySEA' + Processing frames: 0%| | 0/100 [00:00 wfdata torch.Size([1, 100, 481, 481, 1]) +failed to import pySEA: No module named 'pySEA' + Processing frames: 0%| | 0/100 [00:00 wfdata torch.Size([1, 100, 1495, 1295, 1]) -Traceback (most recent call last): - File "/home/qwe/PySlice/TWP/tests/10_midgley.py", line 54, in - tacaw = TACAWData(exitwaves,keep_complex=True) - ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - File "/home/qwe/PySlice/TWP/tests/../src/pyslice/postprocessing/tacaw_data.py", line 87, in __init__ - self._fft_from_wf_data(layer_index) - File "/home/qwe/PySlice/TWP/tests/../src/pyslice/postprocessing/tacaw_data.py", line 252, in _fft_from_wf_data - wf_fft = fftshift(wf_fft, axes=1) - ^^^^^^^^^^^^^^^^^^^^^^^^ - File "/home/qwe/PySlice/TWP/tests/../src/pyslice/backend.py", line 148, in fftshift - return xp.fft.fftshift(k,**kwargs) - ^^^^^^^^^^^^^^^^^^^^^^^^^^^ - File "/home/qwe/.conda/envs/pyslice_env/lib/python3.12/site-packages/torch/utils/_device.py", line 104, in __torch_function__ - return func(*args, **kwargs) - ^^^^^^^^^^^^^^^^^^^^^ -torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 2.89 GiB. GPU 0 has a total capacity of 15.77 GiB of which 937.25 MiB is free. Including non-PyTorch memory, this process has 14.85 GiB memory in use. Of the allocated memory 11.61 GiB is allocated by PyTorch, and 2.87 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation. See documentation for Memory Management (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables) +failed to import pySEA: No module named 'pySEA' + Processing frames: 0%| | 0/100 [00:00