zoomy-core 0.1.1__py3-none-any.whl → 0.1.2__py3-none-any.whl

zoomy_core/misc/misc.py

@@ -0,0 +1,216 @@
+import os
+import numpy as np
+
+# import scipy.interpolate as interp
+# from functools import wraps
+
+from attr import define
+from typing import Callable, Optional, Any
+from types import SimpleNamespace
+
+from sympy import MatrixSymbol
+from sympy import MutableDenseNDimArray as ZArray
+
+from library.zoomy_core.misc.custom_types import FArray
+from library.zoomy_core.misc.static_class import register_static_pytree
+from library.zoomy_core.misc.logger_config import logger
+
+
+
+
+
+
+@register_static_pytree
+@define(slots=True, frozen=False, kw_only=True)
+class Zstruct(SimpleNamespace):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+    def __getitem__(self, key):
+        return self.values()[key]
+
+    def length(self):
+        return len(self.values())
+
+    def get_list(self, recursive: bool = True):
+        if recursive:
+            output = []
+            for item in self.values():
+                if hasattr(item, 'get_list'):
+                    # If item is a Zstruct or similar, call get_list recursively
+                    output.append(item.get_list(recursive=True))
+                else:
+                    output.append(item)
+            return output
+        else:
+            return self.values()
+        
+    def as_dict(self, recursive: bool = True):
+        if recursive:
+            output = {}
+            for key, value in self.items():
+                if hasattr(value, 'as_dict'):
+                    # If value is a Zstruct or similar, call as_dict recursively
+                    output[key] = value.as_dict(recursive=True)
+                else:
+                    output[key] = value
+            return output
+        else:
+            return self.__dict__
+
+    
+    def items(self, resursive: bool = False):
+        return self.as_dict(recursive=resursive).items()
+    
+    def keys(self):
+        return list(self.as_dict(recursive=False).keys())
+    
+    def values(self):
+        return list(self.as_dict(recursive=False).values())
+    
+    def contains(self, key):
+        if self.as_dict(recursive=False).get(key) is not None:
+            return True
+        return False
+    
+    def update(self, zstruct, recursive: bool = True):
+        """
+        Update the current Zstruct with another Zstruct or dictionary.
+        """
+        if not isinstance(zstruct, Zstruct):
+            raise TypeError("zstruct must be a Zstruct or a dictionary.")
+        
+        if recursive:
+            # Update each attribute recursively
+            for key, value in zstruct.as_dict(recursive=False).items():
+                if hasattr(self, key):
+                    current_value = getattr(self, key)
+                    if isinstance(current_value, Zstruct) and isinstance(value, Zstruct):
+                        # If both are Zstructs, update recursively
+                        current_value.update(value, recursive=True)
+                    else:
+                        setattr(self, key, value)
+                else:
+                    setattr(self, key, value)
+        else:
+            # Update only the top-level attributes
+            for key, value in zstruct.as_dict(recursive=False).items():
+                setattr(self, key, value)
+                
+    @classmethod
+    def from_dict(cls, d):
+        """
+        Create a Zstruct recursively from a dictionary.
+        
+        Args:
+            d (dict): Dictionary.
+        
+        Returns:
+            Zstruct: An instance of Zstruct.
+        """
+        if not isinstance(d, dict):
+            raise TypeError("Input must be a dictionary.")
+        
+        # Convert the dictionary to a Zstruct
+        for k, v in d.items():
+            if isinstance(v, dict):
+                d[k] = Zstruct.from_dict(v)     
+        return cls(**d)
+
+    
+@register_static_pytree
+@define(slots=True, frozen=False, kw_only=True)
+class Settings(Zstruct):
+    """
+    Settings class for the application.
+    
+    Args: 
+        **kwargs: Arbitrary keyword arguments to set as attributes.
+        
+    Returns:
+        An `IterableNamespace` instance.
+    """
+    
+    def __init__(self, **kwargs):
+        # assert that kwargs constains name
+        if 'output' not in kwargs or not isinstance(kwargs['output'], Zstruct):
+            logger.warning("No 'output' Zstruct found in Settings. Default: Zstruct(directory='output', filename='simulation', clean_directory=False)")
+            kwargs['output'] = Zstruct(directory='output', filename='simulation', clean_directory=True)
+        output = kwargs['output']
+        if not output.contains('directory'):
+            logger.warning("No 'directory' attribute found in output Zstruct. Default: 'output'")
+            kwargs['output'] = Zstruct(directory='output', **output.as_dict())
+        if not output.contains('filename'):
+            logger.warning("No 'filename' attribute found in output Zstruct. Default: 'simulation'")
+            kwargs['output'] = Zstruct(filename='simulation', **output.as_dict())
+        if not output.contains('clean_directory'):
+            logger.warning("No 'clean_directory' attribute found in output Zstruct. Default: False")
+            kwargs['output'] = Zstruct(clean_directory=False, **output.as_dict())
+        super().__init__(**kwargs)
+        
+    @classmethod
+    def default(cls):
+        """
+        Returns a default Settings instance.
+        """
+        return cls(
+            output=Zstruct(directory='output', filename='simulation', clean_directory=False)
+        )
+    
+
+def compute_transverse_direction(normal):
+    dim = normal.shape[0]
+    if dim == 1:
+        return np.zeros_like(normal)
+    elif dim == 2:
+        transverse = np.zeros((2), dtype=float)
+        transverse[0] = -normal[1]
+        transverse[1] = normal[0]
+        return transverse
+    elif dim == 3:
+        cartesian_x = np.array([1, 0, 0], dtype=float)
+        transverse = np.cross(normal, cartesian_x)
+        return transverse
+    else:
+        assert False
+
+
+def extract_momentum_fields_as_vectors(Q, momentum_fields, dim):
+    num_fields = len(momentum_fields)
+    num_momentum_eqns = int(num_fields / dim)
+    Qnew = np.empty((num_momentum_eqns, dim))
+    for i_eq in range(num_momentum_eqns):
+        for i_dim in range(dim):
+            Qnew[i_eq, i_dim] = Q[momentum_fields[i_dim * num_momentum_eqns + i_eq]]
+    return Qnew
+
+
+def projection_in_normal_and_transverse_direction(Q, momentum_fields, normal):
+    dim = normal.shape[0]
+    transverse_directions = compute_transverse_direction(normal)
+    Q_momentum_eqns = extract_momentum_fields_as_vectors(Q, momentum_fields, dim)
+    Q_normal = np.zeros((Q_momentum_eqns.shape[0]), dtype=float)
+    Q_transverse = np.zeros((Q_momentum_eqns.shape[0]), dtype=float)
+    for d in range(dim):
+        Q_normal += Q_momentum_eqns[:, d] * normal[d]
+        Q_transverse += Q_momentum_eqns[:, d] * transverse_directions[d]
+    return Q_normal, Q_transverse
+
+
+def projection_in_x_y_direction(Qn, Qt, normal):
+    dim = normal.shape[0]
+    num_momentum_fields = Qn.shape[0]
+    transverse_directions = compute_transverse_direction(normal)
+    Q = np.empty((num_momentum_fields * dim), dtype=float)
+    for i in range(num_momentum_fields):
+        for d in range(dim):
+            Q[i + d * num_momentum_fields] = (
+                Qn[i] * normal[d] + Qt[i] * transverse_directions[d]
+            )
+    return Q
+
+
+def project_in_x_y_and_recreate_Q(Qn, Qt, Qorig, momentum_eqns, normal):
+    Qnew = np.array(Qorig)
+    Qnew[momentum_eqns] = projection_in_x_y_direction(Qn, Qt, normal)
+    return Qnew

zoomy_core/misc/static_class.py

@@ -0,0 +1,94 @@
+
+import attrs
+from functools import partial
+from typing import Any, Tuple, Type
+import attr
+
+try:
+    import jax
+    import jax.numpy as jnp
+    _HAVE_JAX = True
+except ImportError:
+    _HAVE_JAX = False
+
+
+def register_static_pytree(cls: Type[Any]) -> Type[Any]:
+    """
+    Class decorator that registers the class as a JAX pytree node,
+    treating all member variables as static.
+
+    Parameters:
+    cls (Type[Any]): The class to register.
+
+    Returns:
+    Type[Any]: The registered class.
+    """
+    if not _HAVE_JAX:
+        # no-op decorator
+        return cls
+    if not attrs.has(cls):
+        raise TypeError(
+            "register_static_pytree can only be applied to classes decorated with @attrs.define or @attr.s"
+        )
+
+    # Extract field names from attrs
+    field_names = [field.name for field in attrs.fields(cls)]
+
+    # Define the flatten function
+    def flatten(instance: Any) -> Tuple[Tuple, Tuple]:
+        aux_data = tuple(getattr(instance, name) for name in field_names)
+        children = ()  # No dynamic children since all are static
+        return children, aux_data
+
+    # Define the unflatten function
+    def unflatten(aux_data: Tuple, children: Tuple) -> Any:
+        return cls(*aux_data)
+
+    # Register the class as a pytree node with JAX
+    jax.tree_util.register_pytree_node(cls, flatten, unflatten)
+
+    return cls
+
+
+# 2. Define the Mesh class using attrs and the decorator
+@register_static_pytree
+@attrs.define(frozen=True)
+class Mesh:
+    x: jnp.ndarray
+    y: jnp.ndarray
+    # Use factory for mutable default fields
+    z: jnp.ndarray = attr.field(factory=lambda: jnp.array([0.0]))
+    # Example additional fields
+    w: jnp.ndarray = attr.field(factory=lambda: jnp.array([[1.0, 2.0], [3.0, 4.0]]))
+    description: str = "Default Mesh Description"
+
+
+# 3. Define the SpaceOperator class with JAX's jit
+class SpaceOperator:
+    @partial(jax.jit, static_argnums=(0,))  # Marks 'self' as static
+    def solve(self, q: jnp.ndarray, mesh: Mesh) -> jnp.ndarray:
+        """
+        Example operation: Element-wise multiplication of mesh.x with q,
+        then add mesh.y for demonstration.
+        """
+        return mesh.x * q + mesh.y
+
+
+if __name__ == "__main__":
+    # Create mesh data using JAX arrays
+    x = jnp.linspace(0, 1, 10)
+    y = jnp.linspace(1, 2, 10)
+    # Initialize Mesh with x and y; z and w use default_factory
+    mesh = Mesh(x, y)
+
+    # Initialize the operator
+    space_op = SpaceOperator()
+
+    # Define the input Q as a JAX array
+    Q = jnp.linspace(0, 1, 10)
+
+    # Use the solve method
+    Q_result = space_op.solve(Q, mesh)
+
+    # Print the result
+    print("Q_result:", Q_result)

zoomy_core/model/analysis.py

@@ -0,0 +1,147 @@
+from sympy import Matrix, diff, exp, I, linear_eq_to_matrix, solve , Eq, zeros, simplify, nsimplify, latex,  symbols, Function, together, Symbol
+from IPython.display import display, Latex
+
+
+class ModelAnalyser():
+    def __init__(self, model):
+        self.model = model
+        self.t = model.time
+        x, y, z = model.position
+        self.x = x
+        self.y = y
+        self.z = z
+        self.equations = None
+        self.plane_wave_symbols = []
+        
+
+    def get_equations(self):
+        return self.equations
+
+    def print_equations(self):          
+        latex_lines = " \\\\\n".join([f"& {latex(eq)}" for eq in self.equations])
+        latex_block = r"$$\begin{align*}" + "\n" + latex_lines + r"\end{align*}$$"
+        display(Latex(latex_block))
+            
+    def get_time_space(self):
+        x, y, z = self.model.position
+        t = self.model.time
+        return t, x, y, z
+    
+    def _get_omega_k(self):
+        omega, kx, ky, kz = symbols('omega k_x k_y k_z')
+        return omega, kx, ky, kz
+    
+    def _get_exponential(self):
+        omega, kx, ky, kz = self._get_omega_k()
+        t, x, y, z = self.get_time_space()
+        exponential = exp(I * (kx * x + ky * y + kz * z - omega * t))
+        return exponential
+    
+    def get_eps(self):
+        eps = symbols('eps')
+        return eps
+    
+    def create_functions_from_list(self, names):
+        t, x, y, z = self.get_time_space()
+        return [Function(name)(t, x, y, z) for name in names]
+    
+    def delete_equations(self, indices):
+        self.equations = [self.equations[i] for i in range(len(self.equations)) if i not in indices]
+
+
+    def solve_for_constraints(self, list_of_selected_equations, list_of_variables):
+        equations = self.equations
+        sol = solve([equations[i] for i in list_of_selected_equations], list_of_variables)
+        equations = [eq.xreplace(sol).doit() for eq in equations]
+        # delete used equations from equation system
+        equations = [equations[i] for i in range(len(equations)) if i not in list_of_selected_equations]
+        self.equations = equations
+        return sol
+
+    def insert_plane_wave_ansatz(self, functions_to_replace):
+        exponential = self._get_exponential()
+        f_bar_dict = {}
+        for f in functions_to_replace:
+            # Create the base name (e.g., 'f0')
+            f_name = str(f.func)  # Get the function name (e.g., 'f0')
+            
+            # Create a new symbol representing \bar{f0}
+            f_bar = Symbol(r'\bar{' + f_name + '}')
+            f_bar_dict[f] =  f_bar * exponential
+            self.plane_wave_symbols.append(f_bar)
+        self.equations =  [eq.xreplace(f_bar_dict).doit() for eq in self.equations ]
+        
+    def solve_for_dispersion_relation(self):
+        assert self.equations is not None, "No equations available to solve for dispersion relation."
+        assert self.plane_wave_symbols, "No plane wave symbols available to solve for dispersion relation. Use insert_plane_wave_ansatz first."
+        A, rhs = linear_eq_to_matrix(self.equations, self.plane_wave_symbols)
+        omega, kx, ky, kz = self._get_omega_k()
+        sol = solve(A.det(), omega)
+        return sol
+
+    def remove_exponential(self):
+        exponential = self._get_exponential()
+        equations = self.equations
+        equations = [simplify(Eq(eq.lhs / exponential, eq.rhs / exponential)) for eq in equations]
+        self.equations = equations
+
+    def linearize_system(self, q, qaux, constraints=None):
+        model = self.model
+        t, x, y, z = self.get_time_space()
+        dim = model.dimension
+        X = [x, y, z]
+        
+        Q = Matrix(model.variables.get_list())
+        Qaux = Matrix(model.aux_variables.get_list())
+
+        
+        substitutions = {Q[i]: q[i] for i in range(len(q))}
+        substitutions.update({Qaux[i]: qaux[i] for i in range(len(qaux))})
+        
+
+        A = model.quasilinear_matrix()
+        S = model.residual()
+        if constraints is not None:
+            C = constraints
+        else:
+            C = zeros(0, 1)
+            
+
+        Q = Q.xreplace(substitutions)
+        for d in range(dim):
+            A[d] = A[d].xreplace(substitutions)
+        S = S.xreplace(substitutions)
+        C = C.xreplace(substitutions)
+        
+        C = C.doit()
+            
+            
+        gradQ = Matrix([diff(q[i], X[j]) for i in range(len(q)) for j in range(dim)]).reshape(len(q), dim)
+            
+        AgradQ = A[0] * gradQ[:, 0]
+        for d in range(1, dim):
+            AgradQ += A[d] * gradQ[:, d]
+
+
+        expr = list(Matrix.vstack((diff(q, t) + AgradQ - S) , C))
+        for i in range(len(expr)):
+            expr[i] = nsimplify(expr[i], rational=True)
+        expr = Matrix(expr)
+        eps = self.get_eps()
+        res = expr.copy()
+        for i, e in enumerate(expr):
+            collected = e
+            collected = collected.series(eps, 0, 2).removeO()
+            order_1_term = collected.coeff(eps, 1)
+            res[i] = order_1_term
+            
+        for r in range(res.shape[0]):
+            denom = together(res[r]).as_numer_denom()[1]
+            res[r] *= denom
+            res[r] = simplify(res[r])        
+        
+        linearized_system = [Eq((res[i]),0) for i in range(res.shape[0])]
+
+        
+        self.equations = linearized_system
+

zoomy_core/model/basefunction.py

@@ -0,0 +1,113 @@
+
+import sympy
+import sympy as sp
+from attrs import define, field
+from sympy import lambdify
+
+from library.zoomy_core.misc.misc import Zstruct
+
+def listify(expr):
+    if type(expr) is sp.Piecewise:
+        return expr
+    else:
+        return list(expr)
+
+def vectorize_constant_sympy_expressions(expr, Q, Qaux):
+    """
+    Replace entries in `expr` that are constant w.r.t. Q and Qaux
+    by entry * ones_like(Q[0]) so NumPy/JAX vectorization works.
+    Handles scalars, lists, sympy.Matrix, sympy.Array, and sympy.Piecewise.
+    """
+    symbol_list = set(Q.get_list() + Qaux.get_list())
+    q0 = Q[0]
+    ones_like = sp.Function("ones_like")  # symbolic placeholder
+    zeros_like = sp.Function("zeros_like")  # symbolic placeholder
+
+    # convert matrices to nested lists (Array handles lists better)
+    if isinstance(expr, (sp.MatrixBase, sp.ImmutableDenseMatrix, sp.MutableDenseMatrix)):
+        expr = expr.tolist()
+
+    def vectorize_entry(entry):
+        """Return entry multiplied by ones_like(q0) if it is constant."""
+        # numeric zero
+        if entry == 0:
+            return zeros_like(q0)
+
+        # numeric constant (int, float, Rational, pi, etc.)
+        if getattr(entry, "is_number", False):
+            return entry * ones_like(q0)
+
+        # symbolic constant independent of Q and Qaux
+        if hasattr(entry, "free_symbols") and entry.free_symbols.isdisjoint(symbol_list):
+            return entry * ones_like(q0)
+
+        # otherwise, depends on variables
+        return entry
+
+    def recurse(e):
+        """Recursively handle Array, Matrix, Piecewise, list, or scalar."""
+        # Handle lists (possibly nested)
+        if isinstance(e, list):
+            return [recurse(sub) for sub in e]
+
+        # Handle Matrices
+        if isinstance(e, sp.MatrixBase):
+            return sp.Matrix([[recurse(sub) for sub in row] for row in e.tolist()])
+
+        # Handle Arrays (any rank)
+        if isinstance(e, sp.Array):
+            return sp.Array([recurse(sub) for sub in e])
+
+        # Handle Piecewise
+        if isinstance(e, sp.Piecewise):
+            # Recursively process all (expr, cond) pairs
+            new_args = []
+            for expr_i, cond_i in e.args:
+                new_expr = recurse(expr_i)
+                new_args.append((new_expr, cond_i))
+            return sp.Piecewise(*new_args)
+
+        # Scalar or atomic expression
+        return vectorize_entry(e)
+
+    # Recurse and then normalize into an N-dimensional array (if possible)
+    result = recurse(expr)
+
+    # Convert to Array if possible (this ensures uniform output type)
+    if isinstance(result, list):
+        try:
+            return sp.Array(result)
+        except Exception:
+            # fall back if shapes inconsistent
+            return result
+    return result
+
+
+
+
+
+@define(frozen=True, slots=True, kw_only=True)
+class Function:
+    """
+    Generic (virtual) function implementation.
+    """
+
+    name: str = field(default="Function")
+    args: Zstruct = field(default=Zstruct())
+    definition = field(default=sympy.zeros(1, 1))
+
+    def __call__(self):
+        """Allow calling the instance to get its symbolic definition."""
+        return self.definition
+
+    def lambdify(self, modules=None):
+        """Return a lambdified version of the function."""
+
+        func = lambdify(
+            self.args.get_list(),
+            vectorize_constant_sympy_expressions(
+                listify(self.definition), self.args.variables, self.args.aux_variables
+            ),
+            modules=modules,
+        )
+        return func