%matplotlib inline
%config Completer.use_jedi = False
import os
from functools import partial
from pathlib import Path
import h5py
import matplotlib.pyplot as plt
import numpy as np
from salvus.mesh import mesh_block
import salvus.mesh.unstructured_mesh as um
import salvus.toolbox.toolbox as st
from salvus.flow import api
from salvus.flow import simple_config as sc
SALVUS_FLOW_SITE_NAME = os.environ.get("SITE_NAME", "local")
ranks = 2
n = 25
x = 1.0
rho = 1000.0
mesh = mesh_block.generators.cartesian.rectangle_2d(
nelem_x=n, nelem_y=n, max_x=x, max_y=x
).get_unstructured_mesh()
mesh.find_side_sets()
mesh_order_1 = mesh.copy()
mesh_order_1.change_tensor_order(1)
mesh_order_4 = mesh.copy()
mesh_order_4.change_tensor_order(4)
axis = np.random.choice([0, 1])
coords_order_1 = mesh_order_1.get_element_nodes()[:, :, axis].flatten()
coords_order_4 = mesh_order_4.get_element_nodes()[:, :, axis].flatten()
def li(x: np.ndarray, min_val, max_val) -> np.ndarray:
"""Linearly interpolate a parameter over a coordinate array.
Parameters
----------
x : np.ndarray
Flattened coordinate array (either x or y).
min_val : float
Value of the parameter at the minimum coordinate.
max_val : float
Value of the parameter at the maximum coordinate.
Returns
-------
np.ndarray
An array of values linearly interpolated over the coordinate range.
"""
m = (max_val - min_val) / (np.max(x) - np.min(x))
return m * x + min_val
vp0, vp1 = 1000, 2000
m10 = m11 = 1 / rho
m00 = m10 / vp0**2
m01 = m11 / vp1**2
m0 = (np.apply_along_axis(li, 0, coords_order_1, m00, m01)).reshape(
mesh_order_1.nelem, mesh_order_1.nodes_per_element
) * 0.5
m1 = (np.apply_along_axis(li, 0, coords_order_1, m10, m11)).reshape(
mesh_order_1.nelem, mesh_order_1.nodes_per_element
) * 0.5
mesh_order_1.attach_field("M0", m0)
mesh_order_1.attach_field("M1", m1)
mesh_order_1.attach_field("fluid", np.ones(mesh.nelem))
mesh_order_1
<salvus.mesh.unstructured_mesh.UnstructuredMesh at 0x7f8f2d35b550>
m0 = (np.apply_along_axis(li, 0, coords_order_4, m00, m01)).reshape(
mesh_order_4.nelem, mesh_order_4.nodes_per_element
) * 0.5
m1 = (np.apply_along_axis(li, 0, coords_order_4, m10, m11)).reshape(
mesh_order_4.nelem, mesh_order_4.nodes_per_element
) * 0.5
mesh_order_4.attach_field("M0", m0)
mesh_order_4.attach_field("M1", m1)
mesh_order_4.attach_field("fluid", np.ones(mesh.nelem))
mesh_order_4