zoomy-core 0.1.14__py3-none-any.whl

model/models/vam.py

@@ -0,0 +1,455 @@
+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
+
+
+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
+
+class VAMHyperbolic(Model):
+    def __init__(
+        self,
+        boundary_conditions,
+        initial_conditions,
+        dimension=1,
+        fields=6,
+        aux_variables=['hw2', 'hp0', 'hp1', 'dbdx', 'dhdx', 'dhp0dx', 'dhp1dx'],
+        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)])
+        hw2 = self.aux_variables.hw2
+        h = self.variables[0]
+        hu0 = self.variables[1]
+        hu1 = self.variables[2]
+        hw0 = self.variables[3]
+        hw1 = self.variables[4]
+        param = self.parameters
+
+        u0 = hu0 / h
+        u1 = hu1 / h
+        w0 = hw0 / h
+        w1 = hw1 / h
+        
+        fx[0] = hu0
+        fx[1] = hu0 * u0 + 1/3 * hu1 * u1 * u1
+        fx[2] = hu0 * w0 + 1/3 * hu1 * w1
+        fx[3] = 2*hu0 * u1
+        fx[4] = hu0 * w1 + u1 * (hw0 + 2/5*hw2)
+        
+        return [fx]
+
+    def nonconservative_matrix(self):
+        nc = Matrix([[0 for i in range(self.n_variables)] for j in range(self.n_variables)])
+
+        hw2 = self.aux_variables.hw2
+        h = self.variables[0]
+        hu0 = self.variables[1]
+        hu1 = self.variables[2]
+        hw0 = self.variables[3]
+        hw1 = self.variables[4]
+        param = self.parameters
+
+        u0 = hu0 / h
+        u1 = hu1 / h
+        w0 = hw0 / h
+        w1 = hw1 / h
+        w2 = hw2 / h
+
+        nc[1, 0] = param.g * h
+        nc[3, 2] = u0
+        nc[4, 2] = - 1/5 * w2 + w0
+        nc[1, 5] = param.g * h
+        return [-nc]
+    
+    def eigenvalues(self):
+        ev = Matrix([0 for i in range(self.n_variables)])
+        h = self.variables[0]
+        hu0 = self.variables[1]
+        hu1 = self.variables[2]
+        param = self.parameters
+
+        u0 = hu0 / h
+        u1 = hu1 / h
+
+        ev[0] = u0
+        ev[1] = u0 + 1/sqrt(3) * u1
+        ev[2] = u0 - 1/sqrt(3) * u1
+        ev[3] = u0 + sqrt(param.g * h + u1**2)
+        ev[4] = u0 - sqrt(param.g * h + u1**2)
+        ev[5] = 0
+        
+        return ev
+
+    def source_implicit(self):
+        R = Matrix([0 for i in range(self.n_variables)])
+        hw2 = self.aux_variables.hw2
+        h = self.variables[0]
+        hu0 = self.variables[1]
+        hu1 = self.variables[2]
+        hw0 = self.variables[3]
+        hw1 = self.variables[4]
+        b = self.variables[5]
+        param = self.parameters
+
+        u0 = hu0 / h
+        u1 = hu1 / h
+        w0 = hw0 / h
+        w1 = hw1 / h
+        w2 = hw2 /h  
+
+
+        p0 = self.aux_variables.hp0/h
+        p1 = self.aux_variables.hp1/h
+        dbdx = self.aux_variables.dbdx
+        dhdx = self.aux_variables.dhdx
+        dhp0dx = self.aux_variables.dhp0dx
+        dhp1dx = self.aux_variables.dhp1dx
+
+        R[0] = 0.
+        R[1] = dhp0dx + 2 * p1 * dbdx 
+        R[2] = -2*p1
+        R[3] = dhp1dx - (3*p0 - p1)*dhdx  -6*(p0-p1)*dbdx
+        R[4] = 6*(p0-p1)
+        R[5] = 0.
+        return R
+        
+
+class VAMPoisson(Model):
+    def __init__(
+        self,
+        boundary_conditions,
+        initial_conditions,
+        dimension=1,
+        fields=['hp0', 'hp1'],
+        aux_variables=['h', 'hu0', 'hu1', 'hw0', 'hw1' ,'b', 'hw2', 'dbdx', 'ddbdxx', 'dhdx', 'ddhdxx', 'du0dx', 'du1dx', 'dhp0dx', 'ddhp0dxx', 'dhp1dx', 'ddhp1dxx', 'dt', 'd4hp0dx4', 'd4hp1dx4'],
+        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 source_implicit(self):
+        R = Matrix([0 for i in range(self.n_variables)])
+
+        h = self.aux_variables.h
+        #p0 = self.variables.p0/h
+        #p1 = self.variables.p1/h
+        dt = self.aux_variables.dt
+
+        dbdx   = self.aux_variables.dbdx
+        ddbdxx = self.aux_variables.ddbdxx
+        dhdx   = self.aux_variables.dhdx
+        ddhdxx = self.aux_variables.ddhdxx
+        #dp0dx   = self.aux_variables.dp0dx
+        #ddp0dxx = self.aux_variables.ddp0dxx
+        #dp1dx   = self.aux_variables.dp1dx
+        #ddp1dxx   = self.aux_variables.ddp1dxx
+        du0dx = self.aux_variables.du0dx
+        du1dx = self.aux_variables.du1dx
+
+        u0 = self.aux_variables.hu0/h
+        u1 = self.aux_variables.hu1/h
+        w0 = self.aux_variables.hw0/h
+        w1 = self.aux_variables.hw1/h
+        
+        hp0 = self.variables.hp0
+        hp1 = self.variables.hp1
+        dhp0dx   = self.aux_variables.dhp0dx
+        ddhp0dxx = self.aux_variables.ddhp0dxx
+        dhp1dx   = self.aux_variables.dhp1dx
+        ddhp1dxx   = self.aux_variables.ddhp1dxx
+        d4hp0dx4   = self.aux_variables.d4hp0dx4
+        d4hp1dx4   = self.aux_variables.d4hp1dx4
+
+
+        
+
+        delta = 0.
+        #I1 = 0.666666666666667*dt*dp0dx - 2*(-dt*(h*ddp0dxx + p0*dhdx + 2*p1*dbdx) + h*dp1dx)*dbdx/h + 2*(-dt*(-(3*p0 - p1)*dhdx - (6*p0 - 6*p1)*dbdx + h*dp0dx + p1*dhdx) + h*u1)/h + 0.333333333333333*(2*dt*p1 + h*u0)*dhdx/h + (-(-dt*(h*ddp0dxx + p0*dhdx + 2*p1*dbdx) + h*dp1dx)*dhdx/h**2 + (-dt*(h*du1dx + p0*ddhdxx + 2*p1*ddbdxx + 2*dbdx*dp0dx + 2*dhdx*ddp0dxx) + h*dhdx + dp1dx*dhdx)/h)*h + 0.333333333333333*h*du0dx + 0.333333333333333*u0*dhdx + delta * ddp0dxx
+        #I2 = -2*(-dt*(6*p0 - 6*p1) + h*w0)/h + 2*(2*dt*p1 + h*u0)*dbdx/h + (2*dt*p1 + h*u0)*dhdx/h + (-(-dt*(h*ddp0dxx + p0*dhdx + 2*p1*dbdx) + h*dp1dx)*dhdx/h**2 + (-dt*(h*du1dx + p0*ddhdxx + 2*p1*ddbdxx + 2*dbdx*dp0dx + 2*dhdx*ddp0dxx) + h*dhdx + dp1dx*dhdx)/h)*h + delta * ddp1dxx
+        I1 = 0.666666666666667*dt*dhp1dx/h - 0.666666666666667*dt*hp1*dhdx/h**2 - 2*(-dt*(dhp0dx + 2*hp1*dbdx/h) + h*u0)*dbdx/h + 2*(-dt*(-(3*hp0/h - hp1/h)*dhdx - (6*hp0/h - 6*hp1/h)*dbdx + dhp1dx) + h*w0)/h + 0.333333333333333*(2*dt*hp1/h + h*u1)*dhdx/h + (-(-dt*(dhp0dx + 2*hp1*dbdx/h) + h*u0)*dhdx/h**2 + (-dt*(ddhp0dxx + 2*hp1*ddbdxx/h + 2*dbdx*dhp1dx/h - 2*hp1*dbdx*dhdx/h**2) + h*du0dx + u0*dhdx)/h)*h + 0.333333333333333*h*du1dx + 0.333333333333333*u1*dhdx + delta *ddhp0dxx
+        I2 = -2*(-dt*(6*hp0/h - 6*hp1/h) + h*w1)/h + 2*(2*dt*hp1/h + h*u1)*dbdx/h + (2*dt*hp1/h + h*u1)*dhdx/h + (-(-dt*(dhp0dx + 2*hp1*dbdx/h) + h*u0)*dhdx/h**2 + (-dt*(ddhp0dxx + 2*hp1*ddbdxx/h + 2*dbdx*dhp1dx/h - 2*hp1*dbdx*dhdx/h**2) + h*du0dx + u0*dhdx)/h)*h + delta *ddhp1dxx
+        R[0] = I1 +I2
+        R[1] = -I2 + I1
+
+        return R
+    
+    def eigenvalues(self):
+        ev = Matrix([0 for i in range(self.n_variables)])
+        return ev
+
+
+class VAMPoissonFull(Model):
+    def __init__(
+        self,
+        boundary_conditions,
+        initial_conditions,
+        dimension=1,
+        fields=8,
+        aux_variables=['dhdt', 'dhu0dt', 'dhu1dt', 'dhw0dt', 'dhw1dt', 'dhdx', 'dhu0dx', 'dhu1dx', 'dhw0dx', 'dhw1dx', 'dhp0dx', 'dhp1dx', 'dbdx', 'hw2', 'ddp0dxx', 'ddp1dxx', 'du0dx', 'du1dx'],
+        parameters={},
+        _default_parameters={"g": 1},
+        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 source_implicit(self):
+        R = Matrix([0 for i in range(self.n_variables)])
+        hw2 = self.aux_variables.hw2
+        h = self.variables[0]
+        hu0 = self.variables[1]
+        hu1 = self.variables[2]
+        hw0 = self.variables[3]
+        hw1 = self.variables[4]
+        b = self.variables[5]
+        p0 = self.variables[6]/h
+        p1 = self.variables[7]/h
+        param = self.parameters
+
+        u0 = hu0 / h
+        u1 = hu1 / h
+        w0 = hw0 / h
+        w1 = hw1 / h
+        w2 = hw2 /h  
+
+
+        dhdt   = self.aux_variables.dhdt   
+        dhu0dt = self.aux_variables.dhu0dt 
+        dhu1dt = self.aux_variables.dhu1dt 
+        dhw0dt = self.aux_variables.dhw0dt 
+        dhw1dt = self.aux_variables.dhw1dt 
+        dhdx   = self.aux_variables.dhdx   
+        dhu0dx = self.aux_variables.dhu0dx 
+        dhu1dx = self.aux_variables.dhu1dx 
+        dhw0dx = self.aux_variables.dhw0dx 
+        dhw1dx = self.aux_variables.dhw1dx 
+        dhp0dx = self.aux_variables.dhp0dx 
+        dhp1dx = self.aux_variables.dhp1dx 
+        dbdx   = self.aux_variables.dbdx   
+        du0dx  = self.aux_variables.du0dx
+        ddp0dxx = self.aux_variables.ddp0dxx
+        ddp1dxx = self.aux_variables.ddp1dxx
+
+        R[0] = dhdt 
+        R[1] = dhu0dt + dhp0dx + 2 * p1 * dbdx 
+        R[2] = dhu1dt -2*p1
+        R[3] = dhw0dt + dhp1dx - (3*p0 - p1)*dhdx  -6*(p0-p1)*dbdx
+        R[4] = 6*(p0-p1)
+        R[5] = 0.
+        I1 = h*du0dx + 1/3 * dhu1dx + 1/3 * u1 * dhdx + 2*(w0 - u0 * dbdx)
+        I2 = h * du0dx + u1*dhdx + 2*(u1 * dbdx - w1)
+        R[6] = I1 + I2
+        R[7] = I1 - I2
+        
+        delta = 0.0
+        R[6] +=  delta * (ddp0dxx )
+        R[7] +=  delta * (ddp1dxx)
+
+        return R
+    
+    def eigenvalues(self):
+        ev = Matrix([0 for i in range(self.n_variables)])
+        return ev
+
+class VAMFullImplicit(Model):
+    def __init__(
+        self,
+        boundary_conditions,
+        initial_conditions,
+        dimension=1,
+        fields=8,
+        aux_variables=[ 'hw2', 'dhdt', 'dhu0dt', 'dhu1dt', 'dhw0dt', 'dhw1dt', 'dhdx', 'dhu0dx', 'dhu1dx', 'dhw0dx', 'dhw1dx', 'dhp0dx', 'dhp1dx', 'dbdx','ddhp0dxx', 'ddhp1dxx', 'du0dx', 'du1dx'],
+        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 source_implicit(self):
+        R = Matrix([0 for i in range(self.n_variables)])
+        hw2 = self.aux_variables.hw2
+        h = self.variables[0]
+        hu0 = self.variables[1]
+        hu1 = self.variables[2]
+        hw0 = self.variables[3]
+        hw1 = self.variables[4]
+        b = self.variables[5]
+        p0 = self.variables[6]/h
+        p1 = self.variables[7]/h
+        param = self.parameters
+
+        u0 = hu0 / h
+        u1 = hu1 / h
+        w0 = hw0 / h
+        w1 = hw1 / h
+        w2 = hw2 /h  
+
+
+        dhdt   = self.aux_variables.dhdt   
+        dhu0dt = self.aux_variables.dhu0dt 
+        dhu1dt = self.aux_variables.dhu1dt 
+        dhw0dt = self.aux_variables.dhw0dt 
+        dhw1dt = self.aux_variables.dhw1dt 
+        dhdx   = self.aux_variables.dhdx   
+        dhu0dx = self.aux_variables.dhu0dx 
+        dhu1dx = self.aux_variables.dhu1dx 
+        dhw0dx = self.aux_variables.dhw0dx 
+        dhw1dx = self.aux_variables.dhw1dx 
+        dhp0dx = self.aux_variables.dhp0dx 
+        dhp1dx = self.aux_variables.dhp1dx 
+        dbdx   = self.aux_variables.dbdx   
+        du0dx  = self.aux_variables.du0dx
+        ddhp0dxx = self.aux_variables.ddhp0dxx
+        ddhp1dxx = self.aux_variables.ddhp1dxx
+
+        R[0] = dhdt 
+        R[1] = dhu0dt + dhp0dx + 2 * p1 * dbdx 
+        R[2] = dhu1dt -2*p1
+        R[3] = dhw0dt + dhp1dx - (3*p0 - p1)*dhdx  -6*(p0-p1)*dbdx
+        R[4] = 6*(p0-p1)
+        R[5] = 0.
+        I1 = h*du0dx + 1/3 * dhu1dx + 1/3 * u1 * dhdx + 2*(w0 - u0 * dbdx)
+        I2 = h * du0dx + u1*dhdx + 2*(u1 * dbdx - w1)
+        R[6] = I1 
+        R[7] = I2
+        
+        delta = 0.1
+        R[6] +=  delta * (ddhp0dxx )
+        R[7] +=  delta * (ddhp1dxx)
+
+        return R
+
+    def flux(self):
+        fx = Matrix([0 for i in range(self.n_variables)])
+        hw2 = self.aux_variables.hw2
+        h = self.variables[0]
+        hu0 = self.variables[1]
+        hu1 = self.variables[2]
+        hw0 = self.variables[3]
+        hw1 = self.variables[4]
+        param = self.parameters
+
+        u0 = hu0 / h
+        u1 = hu1 / h
+        w0 = hw0 / h
+        w1 = hw1 / h
+        
+        fx[0] = hu0
+        fx[1] = hu0 * u0 + 1/3 * hu1 * u1 * u1
+        fx[2] = hu0 * w0 + 1/3 * hu1 * w1
+        fx[3] = 2*hu0 * u1
+        fx[4] = hu0 * w1 + u1 * (hw0 + 2/5*hw2)
+        
+        return [fx]
+
+    def nonconservative_matrix(self):
+        nc = Matrix([[0 for i in range(self.n_variables)] for j in range(self.n_variables)])
+
+        hw2 = self.aux_variables.hw2
+        h = self.variables[0]
+        hu0 = self.variables[1]
+        hu1 = self.variables[2]
+        hw0 = self.variables[3]
+        hw1 = self.variables[4]
+        param = self.parameters
+
+        u0 = hu0 / h
+        u1 = hu1 / h
+        w0 = hw0 / h
+        w1 = hw1 / h
+        w2 = hw2 / h
+
+        nc[1, 0] = param.g * h
+        nc[3, 2] = u0
+        nc[4, 2] = - 1/5 * w2 + w0
+        nc[1, 5] = param.g * h
+        return [-nc]
+    
+    def eigenvalues(self):
+        ev = Matrix([0 for i in range(self.n_variables)])
+        h = self.variables[0]
+        hu0 = self.variables[1]
+        hu1 = self.variables[2]
+        param = self.parameters
+
+        u0 = hu0 / h
+        u1 = hu1 / h
+
+        ev[0] = u0
+        ev[1] = u0 + 1/sqrt(3) * u1
+        ev[2] = u0 - 1/sqrt(3) * u1
+        ev[3] = u0 + sqrt(param.g * h + u1**2)
+        ev[4] = u0 - sqrt(param.g * h + u1**2)
+        ev[5] = 0
+        ev[6] = 0
+        ev[7] = 0
+        
+        return ev
+    
+    
+    

postprocessing/plotting.py

@@ -0,0 +1,244 @@
+import os
+import json
+import meshio
+import numpy as np
+import matplotlib.pyplot as plt
+from matplotlib.collections import PolyCollection
+
+
+# ---------- I/O ----------
+
+
+def read_vtk_or_series(path, index=0, verbose=True):
+    """Read either .vtk or .vtk.series (relative paths handled)."""
+    path = os.path.abspath(path)
+
+    if path.endswith(".vtk"):
+        if verbose:
+            print(f"📁 Reading single VTK file: {path}")
+        return meshio.read(path)
+
+    if path.endswith(".vtk.series"):
+        with open(path, "r") as f:
+            series = json.load(f)
+
+        entries = series.get("files", [])
+        files, times = [], []
+        for e in entries:
+            fname = e.get("name") or e.get("filename") or e.get("file")
+            if fname:
+                files.append(fname)
+                times.append(e.get("time", None))
+
+        if not files:
+            raise ValueError("No valid file entries found in .vtk.series.")
+
+        n_files = len(files)
+        if verbose:
+            print(f"📂 Loaded VTK series: {path}")
+            print(f"   • Available indices: {n_files}")
+            if any(t is not None for t in times):
+                print("   • Time values:", [t for t in times if t is not None])
+            print(f"   • Reading index {index}")
+
+        if not (0 <= index < n_files):
+            raise IndexError(f"Index {index} out of range (0..{n_files - 1})")
+
+        base_dir = os.path.dirname(path)
+        file_path = os.path.join(base_dir, files[index])
+        return meshio.read(file_path)
+
+    raise ValueError("Expected a .vtk or .vtk.series file")
+
+
+# ---------- Field listing ----------
+
+
+def list_available_fields(mesh, print_out=True):
+    """List all available point and cell fields."""
+    fields = {
+        "point_data": list(mesh.point_data.keys()),
+        "cell_data": list(mesh.cell_data_dict.keys()),
+    }
+
+    if print_out:
+        print("\n📊 Available fields:")
+        if fields["point_data"]:
+            print("  • Point data:")
+            for name in fields["point_data"]:
+                arr = mesh.point_data[name]
+                print(f"     - {name} ({arr.shape})")
+        else:
+            print("  • Point data: none")
+
+        if fields["cell_data"]:
+            print("  • Cell data:")
+            for name in fields["cell_data"]:
+                entry = mesh.cell_data_dict[name]
+                types = ", ".join(entry.keys())
+                print(f"     - {name} (types: {types})")
+        else:
+            print("  • Cell data: none")
+        print()
+
+    return fields
+
+
+# ---------- Cell helpers ----------
+
+
+def get_cell_block(mesh, types):
+    for block in mesh.cells:
+        if block.type in types:
+            return block
+    return None
+
+
+def get_cell_field_for_block(mesh, field_name, cell_block):
+    """Get cell-centered data for the specified block."""
+    if cell_block is None:
+        return None
+    ctype = cell_block.type
+    if field_name is not None:
+        field = mesh.cell_data_dict.get(field_name)
+        if field is None:
+            return None
+        return field.get(ctype)
+    for name, mapping in mesh.cell_data_dict.items():
+        if ctype in mapping:
+            return mapping[ctype]
+    return None
+
+
+# ---------- 1D Plotting ----------
+
+
+def plot_1d_mesh(mesh, ax, field_name=None, **_):
+    """
+    Plot a 1D mesh as a line plot on the given Axes.
+    Returns (vmin, vmax) for consistency.
+    """
+    points = mesh.points
+    x = points[:, 0]
+
+    line_block = get_cell_block(mesh, ("line", "line3"))
+    if line_block is not None:
+        connectivity = line_block.data
+        unique_order = []
+        for conn in connectivity:
+            for i in conn:
+                if i not in unique_order:
+                    unique_order.append(i)
+        x = x[unique_order]
+
+    y = None
+    field_data = get_cell_field_for_block(mesh, field_name, line_block)
+    if field_data is not None:
+        cell_centers = np.array(
+            [np.mean(mesh.points[cell, 0]) for cell in line_block.data]
+        )
+        y = np.interp(x, cell_centers, field_data)
+    elif field_name and field_name in mesh.point_data:
+        y = mesh.point_data[field_name]
+    elif mesh.point_data:
+        field_name, y = next(iter(mesh.point_data.items()))
+
+    if y is None:
+        y = np.zeros_like(x)
+        field_name = field_name or "default"
+
+    ax.plot(x, y, "-o", markersize=3)
+    return float(np.min(y)), float(np.max(y))
+
+
+# ---------- 2D Plotting ----------
+
+
+def plot_2d_mesh(
+    mesh,
+    ax,
+    field_name=None,
+    show_legend=True,
+    legend_location="right",
+    cmap="viridis",
+):
+    """
+    Plot a 2D mesh (unstructured) using meshio connectivity.
+    Adds polygons to the provided Axes.
+    Returns (vmin, vmax).
+    """
+    points = mesh.points[:, :2]
+    cell_block = get_cell_block(mesh, ("triangle", "quad", "polygon"))
+    if cell_block is None:
+        raise ValueError("No 2D cells found.")
+
+    connectivity = cell_block.data
+    field_data = get_cell_field_for_block(mesh, field_name, cell_block)
+    if field_data is None:
+        field_data = np.zeros(len(connectivity))
+        field_name = field_name or "default"
+
+    polygons = [points[cell] for cell in connectivity]
+    vmin, vmax = float(field_data.min()), float(field_data.max())
+    if vmin == vmax:
+        vmin, vmax = vmin - 0.5, vmax + 0.5
+
+    coll = PolyCollection(
+        polygons,
+        array=field_data,
+        cmap=cmap,
+        norm=plt.Normalize(vmin=vmin, vmax=vmax),
+        edgecolors="k",
+        linewidths=0.2,
+    )
+    ax.add_collection(coll)
+    ax.autoscale()
+    ax.set_aspect("equal")
+
+    if show_legend:
+        if legend_location == "bottom":
+            cbar = plt.colorbar(
+                coll, ax=ax, orientation="horizontal", fraction=0.05, pad=0.05
+            )
+        else:  # right
+            cbar = plt.colorbar(coll, ax=ax)
+        cbar.set_label(field_name)
+
+    return vmin, vmax
+
+
+# ---------- Auto Dispatch ----------
+
+
+def plot_mesh(
+    mesh,
+    ax,
+    field_name=None,
+    show_legend=True,
+    legend_location="right",
+):
+    """
+    Auto-detect mesh type (1D or 2D) and plot on the provided Axes.
+    Returns (vmin, vmax).
+    """
+    has_2d = get_cell_block(mesh, ("triangle", "quad", "polygon")) is not None
+    has_1d = get_cell_block(mesh, ("line", "line3")) is not None
+
+    if has_2d:
+        return plot_2d_mesh(mesh, ax, field_name, show_legend, legend_location)
+    elif has_1d:
+        return plot_1d_mesh(mesh, ax, field_name)
+    else:
+        raise ValueError(f"Unsupported mesh types: {[c.type for c in mesh.cells]}")
+
+
+# # ---------- Example ----------
+
+# if __name__ == "__main__":
+#     mesh = read_vtk_or_series("outputs/swe/out.vtk.series", index=0)
+#     list_available_fields(mesh)
+
+#     fig, ax = plt.subplots(figsize=(5, 4))
+#     vmin, vmax = plot_mesh(mesh, ax, field_name=None, show_legend=True, legend_location="right")
+#     print(f"Color range: {vmin:.3g} – {vmax:.3g}")
+#     plt.show()