zoomy-core 0.1.14__py3-none-any.whl

model/initial_conditions.py

@@ -0,0 +1,171 @@
+import numpy as np
+from attr import define
+from typing import Callable, Optional
+
+from zoomy_core.misc.custom_types import FArray
+from zoomy_core.mesh.mesh import Mesh
+import zoomy_core.misc.io as io
+import zoomy_core.misc.interpolation as interpolate_mesh
+
+# from zoomy_core.mesh import Mesh2D
+
+
+@define(slots=True, frozen=True)
+class InitialConditions:
+    def apply(self, X, Q):
+        assert False
+        return Q
+
+
+@define(slots=True, frozen=True)
+class Constant(InitialConditions):
+    constants: Callable[[int], FArray] = lambda n_variables: np.array(
+        [1.0] + [0.0 for i in range(n_variables - 1)]
+    )
+
+    def apply(self, X, Q):
+        n_variables = Q.shape[0]
+        for i in range(Q.shape[1]):
+            Q[:, i] = self.constants(n_variables)
+        return Q
+
+
+@define(slots=True, frozen=False)
+class RP(InitialConditions):
+    low: Callable[[int], FArray] = lambda n_variables: np.array(
+        [1.0 * (i == 0) for i in range(n_variables)]
+    )
+    high: Callable[[int], FArray] = lambda n_variables: np.array(
+        [2.0 * (i == 0) for i in range(n_variables)]
+    )
+    jump_position_x: float = 0.0
+
+    def apply(self, X, Q):
+        assert X.shape[1] == Q.shape[1]
+        n_variables = Q.shape[0]
+        for i in range(Q.shape[1]):
+            if X[0, i] < self.jump_position_x:
+                Q[:, i] = self.high(n_variables)
+            else:
+                Q[:, i] = self.low(n_variables)
+        return Q
+
+
+@define(slots=True, frozen=False)
+class RP2d(InitialConditions):
+    low: Callable[[int], FArray] = lambda n_variables: np.array(
+        [1.0 * (i == 0) for i in range(n_variables)]
+    )
+    high: Callable[[int], FArray] = lambda n_variables: np.array(
+        [2.0 * (i == 0) for i in range(n_variables)]
+    )
+    jump_position_x: float = 0.0
+    jump_position_y: float = 0.0
+
+    def apply(self, X, Q):
+        assert X.shape[1] == Q.shape[1]
+        n_variables = Q.shape[0]
+        for i in range(Q.shape[1]):
+            if X[0, i] < self.jump_position_x and X[1, i] < self.jump_position_y:
+                Q[:, i] = self.high(n_variables)
+            else:
+                Q[:, i] = self.low(n_variables)
+        return Q
+
+
+@define(slots=True, frozen=False)
+class RP3d(InitialConditions):
+    low: Callable[[int], FArray] = lambda n_variables: np.array(
+        [1.0 * (i == 0) for i in range(n_variables)]
+    )
+    high: Callable[[int], FArray] = lambda n_variables: np.array(
+        [2.0 * (i == 0) for i in range(n_variables)]
+    )
+    jump_position_x: float = 0.0
+    jump_position_y: float = 0.0
+    jump_position_z: float = 0.0
+
+    def apply(self, X, Q):
+        assert X.shape[1] == Q.shape[1]
+        n_variables = Q.shape[0]
+        for i in range(Q.shape[1]):
+            if (
+                X[0, i] < self.jump_position_x
+                and X[1, i] < self.jump_position_y
+                and X[2, i] < self.jump_position_z
+            ):
+                Q[:, i] = self.high(n_variables)
+            else:
+                Q[:, i] = self.low(n_variables)
+        return Q
+
+
+@define(slots=True, frozen=False)
+class RadialDambreak(InitialConditions):
+    low: Callable[[int], FArray] = lambda n_variables: np.array(
+        [1.0 * (i == 0) for i in range(n_variables)]
+    )
+    high: Callable[[int], FArray] = lambda n_variables: np.array(
+        [2.0 * (i == 0) for i in range(n_variables)]
+    )
+    radius: float = 0.1
+
+    def apply(self, X, Q):
+        dim = X.shape[0]
+        center = np.zeros(dim)
+        for d in range(dim):
+            center[d] = X[d, :].mean()
+        assert X.shape[1] == Q.shape[1]
+        n_variables = Q.shape[0]
+        for i in range(Q.shape[1]):
+            if np.linalg.norm(X[:, i] - center) <= self.radius:
+                Q[:, i] = self.high(n_variables)
+            else:
+                Q[:, i] = self.low(n_variables)
+        return Q
+
+
+@define(slots=True, frozen=True)
+class UserFunction(InitialConditions):
+    function: Optional[Callable[[FArray], FArray]] = None
+
+    def apply(self, X, Q):
+        assert X.shape[1] == Q.shape[1]
+        if self.function is None:
+            self.function = lambda x: np.zeros(Q.shape[1])
+        for i, x in enumerate(X.T):
+            Q[:, i] = self.function(x)
+        return Q
+
+
+# TODO do time interpolation
+@define(slots=True, frozen=True)
+class RestartFromHdf5(InitialConditions):
+    path_to_fields: Optional[str] = None
+    mesh_new: Optional[Mesh] = None
+    mesh_identical: bool = False
+    path_to_old_mesh: Optional[str] = None
+    snapshot: Optional[int] = -1
+    map_fields: Optional[dict] = None
+
+    def apply(self, X, Q):
+        assert self.mesh_new is not None
+        assert self.path_to_fields is not None
+        assert X.shape[0] == Q.shape[0]
+        if self.map_fields is None:
+            map_fields = {i: i for i in range(Q.shape[1])}
+        else:
+            map_fields = self.map_fields
+        mesh = Mesh.from_hdf5(self.path_to_old_mesh)
+        _Q, _Qaux, time = io.load_fields_from_hdf5(
+            self.path_to_fields, i_snapshot=self.snapshot
+        )
+        Q = np.zeros_like(Q)
+        if self.mesh_identical:
+            Q[:, list(map_fields.values())] = _Q[:, list(map_fields.keys())]
+        else:
+            assert self.path_to_old_mesh is not None
+            Q[:, list(map_fields.values())] = interpolate_mesh.to_new_mesh(
+                _Q, mesh, self.mesh_new
+            )[:, list(map_fields.keys())]
+        return Q

model/model.py

@@ -0,0 +1,65 @@
+import os
+import numpy as np
+from typing import Union, Type
+
+from zoomy_core.model.basemodel import Model
+from zoomy_core.model.models.advection import Advection
+#
+
+# from zoomy_core.model.models.shallow_moments_sediment import *
+import zoomy_core.model.initial_conditions as IC
+import zoomy_core.model.boundary_conditions as BC
+from zoomy_core.mesh.fvm_mesh import Mesh
+from zoomy_core.misc import misc as misc
+
+
+def create_default_mesh_and_model(
+    dimension: int = 1,
+    cls: Type[Model] = Advection,
+    fields: Union[int, list] = 1,
+    aux_variables: Union[int, list] = 0,
+    parameters: Union[int, list, dict] = 0,
+    settings: dict = {},
+):
+    main_dir = misc.get_main_directory()
+
+    assert main_dir != ""
+    ic = IC.Constant()
+
+    bc_tags = ["left", "right", "top", "bottom"][: 2 * dimension]
+    bcs = BC.BoundaryConditions([BC.Wall(tag=tag) for tag in bc_tags])
+    if dimension == 1:
+        mesh = Mesh.create_1d((-1, 1), 10)
+    elif dimension == 2:
+        mesh = Mesh.load_mesh(
+            os.path.join(main_dir, "meshes/quad_2d/mesh_coarse.msh"),
+            "quad",
+            2,
+            bc_tags,
+        )
+    else:
+        assert False
+    model = cls(
+        dimension=dimension,
+        fields=fields,
+        aux_variables=aux_variables,
+        parameters=parameters,
+        boundary_conditions=bcs,
+        initial_conditions=ic,
+        settings=settings,
+    )
+    n_elements = mesh.n_elements
+
+    Q = np.linspace(1, fields * n_elements, fields * n_elements).reshape(
+        n_elements, fields
+    )
+    Qaux = np.zeros((Q.shape[0], model.aux_variables.length()))
+    parameters = model.parameter_values
+    num_normals = mesh.element_n_neighbors
+    normals = np.array(
+        [mesh.element_face_normals[:, i] for i in range(mesh.n_faces_per_element)]
+    )
+
+    model.initial_conditions.apply(Q, mesh.element_center)
+    # model.boundary_conditions.apply()
+    return mesh, model, Q, Qaux, parameters, num_normals, normals

model/models/GN.py

@@ -0,0 +1,70 @@
+import numpy as np
+import numpy.polynomial.legendre as L
+import numpy.polynomial.chebyshev as C
+from scipy.optimize import least_squares as lsq
+import sympy
+from sympy import Matrix, sqrt
+from sympy.abc import x
+
+from sympy import integrate, diff
+from sympy import legendre
+from sympy import lambdify, Rational
+
+
+from zoomy_core.model.basemodel import (
+    register_sympy_attribute,
+    eigenvalue_dict_to_matrix,
+)
+from zoomy_core.model.basemodel import Model
+import zoomy_core.model.initial_conditions as IC
+import zoomy_core.model.boundary_conditions as BC
+
+class GN(Model):
+    def __init__(
+        self,
+        boundary_conditions= None,
+        initial_conditions=None,
+        dimension=1,
+        fields=2,
+        # D = h^3 / 3 * (dt * dx * u + u * dx^2 u - (dx u)^2)
+        aux_variables=['dD_dx'],
+        parameters={},
+        _default_parameters={"g": 9.81},
+        settings={},
+        settings_default={},
+    ):
+        self.variables = register_sympy_attribute(fields, "q")
+        self.n_variables = self.variables.length()
+        super().__init__(
+            dimension=dimension,
+            fields=fields,
+            aux_variables=aux_variables,
+            parameters=parameters,
+            _default_parameters=_default_parameters,
+            boundary_conditions=boundary_conditions,
+            initial_conditions=initial_conditions,
+            settings={**settings_default, **settings},
+        )
+        
+    def flux(self):
+        fx = Matrix([0 for i in range(self.n_variables)])
+        h = self.variables[0]
+        hu = self.variables[1]
+
+        param = self.parameters
+        
+        fx[0] = hu
+        fx[1] = hu**2 / h + 1/2 * param.g * h**2 
+        return [fx]
+
+    
+
+    def source_implicit(self):
+        R = Matrix([0 for i in range(self.n_variables)])
+        dD_dx = self.aux_variables.dD_dx
+
+
+        R[0] = 0
+        R[1] = dD_dx
+        return R
+        

model/models/advection.py

@@ -0,0 +1,53 @@
+import numpy as np
+import os
+# import logging
+
+import sympy
+from sympy import Symbol, Matrix, lambdify, transpose, Abs, sqrt
+
+from sympy import zeros, ones
+
+from attr import define
+from typing import Optional
+from types import SimpleNamespace
+
+from zoomy_core.model.boundary_conditions import BoundaryConditions, Extrapolation
+from zoomy_core.model.initial_conditions import InitialConditions, Constant
+from zoomy_core.misc.custom_types import FArray
+
+# from zoomy_core.misc import vectorize  # type: ignore
+from zoomy_core.model.basemodel import Model
+
+
+class Advection(Model):
+    def flux(self):
+        if self.dimension == 1:
+            F = Matrix([0 for i in range(self.n_variables)])
+            for i_field in range(self.n_variables):
+                F[i_field] = self.variables[i_field] * self.parameters[0]
+            return [F]
+        elif self.dimension == 2:
+            F = Matrix([0 for i in range(self.n_variables)])
+            G = Matrix([0 for i in range(self.n_variables)])
+            for i_field in range(self.n_variables):
+                F[i_field] = self.variables[i_field] * self.parameters[0]
+                G[i_field] = self.variables[i_field] * self.parameters[1]
+            return [F, G]
+        elif self.dimension == 3:
+            F = Matrix([0 for i in range(self.n_variables)])
+            G = Matrix([0 for i in range(self.n_variables)])
+            H = Matrix([0 for i in range(self.n_variables)])
+            for i_field in range(self.n_variables):
+                F[i_field] = self.variables[i_field] * self.parameters[0]
+                G[i_field] = self.variables[i_field] * self.parameters[1]
+                H[i_field] = self.variables[i_field] * self.parameters[2]
+            return [F, G, H]
+        else:
+            assert False
+
+    # def eigenvalues(self):
+    #     assert self.normal.shape[0] == self.parameters.shape[0]
+    #     ev = self.normal[0] * self.parameters[0]
+    #     for d in range(1, self.dimension):
+    #         ev += self.normal[d] * self.parameters[d]
+    #     self.sympy_eigenvalues = Matrix[[ev for i in range(self.n_variables)]]

model/models/basisfunctions.py

@@ -0,0 +1,181 @@
+from copy import deepcopy
+
+import numpy as np
+import sympy
+from sympy import Symbol, bspline_basis_set, diff, integrate, lambdify, legendre
+from sympy.abc import z
+from sympy.functions.special.polynomials import chebyshevu
+
+
+class Basisfunction:
+    name = "Basisfunction"
+    
+    def bounds(self):
+        return [0, 1]
+
+    def basis_definition(self):
+        z = Symbol("z")
+        b = lambda k, z: z**k
+        return [b(k, z) for k in range(self.level + 1)]
+    
+    def weight(self, z):
+        return 1
+    
+    def weight_eval(self, z):
+        z = Symbol("z")
+        f = sympy.lambdify(z, self.weight(z))
+        return f(z)
+
+    def __init__(self, level=0, **kwargs):
+        self.level = level
+        self.basis = self.basis_definition(**kwargs)
+
+    def get(self, k):
+        return self.basis[k]
+
+    def eval(self, k, _z):
+        return self.get(k).subs(z, _z)
+    
+    def eval_psi(self, k, _z):
+        z = sympy.Symbol('z')
+        psi = sympy.integrate(self.get(k), (z, self.bounds()[0], z))
+        return psi.subs(z, _z)
+
+    def get_lambda(self, k):
+        f = lambdify(z, self.get(k))
+
+        def lam(z):
+            if type(z) == int or type(z) == float:
+                return f(z)
+            elif type(z) == list or type(z) == np.ndarray:
+                return np.array([f(xi) for xi in z])
+            else:
+                assert False
+
+        return lam
+
+    def plot(self, ax):
+        X = np.linspace(self.bounds()[0], self.bounds()[1], 1000)
+        for i in range(len(self.basis)):
+            f = lambdify(z, self.get(i))
+            y = np.array([f(xi) for xi in X])
+            ax.plot(X, y, label=f"basis {i}")
+
+    def reconstruct_velocity_profile(self, alpha, N=100):
+        Z = np.linspace(self.bounds()[0], self.bounds()[1], N)
+        u = np.zeros_like(Z)
+        for i in range(len(self.basis)):
+            b = lambdify(z, self.get(i))
+            u[:] += alpha[i] * b(Z)
+        return u
+    
+    def reconstruct_velocity_profile_at(self, alpha, z):
+        u = 0
+        for i in range(len(self.basis)):
+            b = lambdify(z, self.eval(i, z))
+            u += alpha[i] * b(z)
+        return u
+
+    def reconstruct_alpha(self, velocities, z):
+        n_basis = len(self.basis)
+        alpha = np.zeros(n_basis)
+        for i in range(n_basis):
+            b = lambdify(z, self.eval(i, z))
+            nom = np.trapz(velocities * b(z) * self.weight(z), z)
+            if type(b(z)) == int:
+                den = b(z) ** 2
+            else:
+                den = np.trapz((b(z) * b(z)).reshape(z.shape), z)
+            res = nom / den
+            alpha[i] = res
+        return alpha
+    
+    def project_onto_basis(self, Y):
+        Z = np.linspace(self.bounds()[0], self.bounds()[1], Y.shape[0])
+        n_basis = len(self.basis)
+        alpha = np.zeros(n_basis)
+        z = Symbol("z")
+        for i in range(n_basis):
+            b = lambdify(z, self.eval(i, Z))
+            alpha[i] = np.trapz(Y * b(Z) * self.weight_eval(Z), Z)
+        return alpha
+
+    def get_diff_basis(self):
+        db = [diff(b, z) for i, b in enumerate(self.basis)]
+        return db
+
+
+class Monomials(Basisfunction):
+    name = "Monomials"
+
+
+class Legendre_shifted(Basisfunction):
+    name = "Legendre_shifted"
+
+    def basis_definition(self):
+        z = Symbol("z")
+        b = lambda k, z: legendre(k, 2 * z - 1) * (-1) ** (k)
+        return [b(k, z) for k in range(self.level + 1)]
+    
+class Chebyshevu(Basisfunction):
+    name = "Chebyshevu"
+    
+    def bounds(self):
+        return [-1, 1]
+    
+    def weight(self, z):
+        # do not forget to include the jacobian of the coordinate transformation in the weight
+        return sympy.sqrt(1-z**2)
+
+    def basis_definition(self):
+        z = Symbol("z")
+        b = lambda k, z: sympy.sqrt(2 / sympy.pi) * chebyshevu(k, z)
+        return [b(k, z) for k in range(self.level + 1)]
+    
+class Legendre_DN(Basisfunction):
+    name = "Legendre_DN - satifying no-slip and no-stress. This is a non-SWE basis"
+
+    def bounds(self):
+        return [-1, 1]
+
+    def basis_definition(self):
+        z = Symbol("z")
+        def b(k, z):
+            alpha = sympy.Rational((2*k+3), (k+2)**2)
+            beta = -sympy.Rational((k+1),(k+2))**2
+            return (legendre(k, z) ) + alpha * (legendre(k+1, z) ) + beta * (legendre(k+2, z))
+        #normalizing makes no sence, as b(k, 0) = 0 by construction
+        return [b(k, z) for k in range(self.level + 1)]
+
+
+class Spline(Basisfunction):
+    name = "Spline"
+
+    def basis_definition(self, degree=1, knots=[0, 0, 0.001, 1, 1]):
+        z = Symbol("z")
+        basis = bspline_basis_set(degree, knots, z)
+        return basis
+
+
+class OrthogonalSplineWithConstant(Basisfunction):
+    name = "OrthogonalSplineWithConstant"
+
+    def basis_definition(self, degree=1, knots=[0, 0, 0.5, 1, 1]):
+        z = Symbol("z")
+
+        def prod(u, v):
+            return integrate(u * v, (z, 0, 1))
+
+        basis = bspline_basis_set(degree, knots, z)
+        add_basis = [1]
+        # add_basis = [sympy.Piecewise((0, z<0.1), (1, True))]
+        basis = add_basis + basis[:-1]
+        orth = deepcopy(basis)
+        for i in range(1, len(orth)):
+            for j in range(0, i):
+                orth[i] -= prod(basis[i], orth[j]) / prod(orth[j], orth[j]) * orth[j]
+
+        for i in range(len(orth)):
+            orth[i] /= sympy.sqrt(prod(orth[i], orth[i]))
+
+        return orth