import jax
import jax.numpy as jnp
from abc import abstractmethod
from jax_sn.tools.mesh_tools import compute_face_normal_vectors_from_3d_simplices, _compute_adjacency, _compute_face_permutations, _pad_last_dimension
from typing import Tuple
from .cross_sections import CrossSectionDataSingleGroup
from typing import Tuple
import equinox as eqx
from typing import Union
[docs]
def compute_face_normals(vertices : jnp.ndarray, connectivity : jnp.ndarray, face_template : jnp.ndarray) -> jnp.ndarray:
'''
The canonical face-normal derivation: the single place that knows how ``face_normals`` follows from
``vertices``. Traceable, so it can be re-derived under ``jit``/``grad`` when the vertices move.
Parameters
----------
vertices : jnp.ndarray [n_vertices, n_dim]
connectivity : jnp.ndarray [n_elements, n_vertices_per_element]
face_template : jnp.ndarray [n_faces_per_element, n_vertices_per_face]
Returns
-------
face_normals : jnp.ndarray [n_elements, n_faces_per_element, 3]
'''
return compute_face_normal_vectors_from_3d_simplices(vertices[connectivity], face_template)
[docs]
@jax.jit
def compute_topology(connectivity : jnp.ndarray, face_template : jnp.ndarray):
'''
The mesh adjacency: ``etoe``, ``etof`` and ``face_perm_idx``.
Depends on the connectivity alone, never on the vertices -- which is what lets a deformable
parametrisation move the vertices without redoing this (the expensive, non-differentiable part
of the domain setup).
Parameters
----------
connectivity : jnp.ndarray [n_elements, n_vertices_per_element]
face_template : jnp.ndarray [n_faces_per_element, n_vertices_per_face]
Returns
-------
etoe, etof, face_perm_idx : jnp.ndarray, each [n_elements, n_faces_per_element]
'''
etoe, etof = _compute_adjacency(connectivity, face_template)
face_perm_idx = _compute_face_permutations(connectivity, etoe, etof, face_template)
return etoe, etof, face_perm_idx
@jax.jit
def _compute_domain_setup_from_mesh(mesh : Tuple[jnp.ndarray, jnp.ndarray], face_template : jnp.ndarray):
'''
Computes the setup for a geometric domain from a mesh and a face template (see `jax_sn.tools.mesh_tools._face_template` for an example template).
Computes normals of each face and the element-to-element and element-to-face adjacency matrices.
Parameters
----------
mesh : tuple
A tuple containing vertices and connectivity of the mesh. vertices is a jnp.ndarray of shape [n_vertices, n_dim] and connectivity is a jnp.ndarray of shape [n_elements, n_vertices_per_element].
face_template : jnp.ndarray [n_faces_per_element, n_vertices_per_face]
A 2D array representing the face template.
'''
assert face_template.ndim == 2, "Face template must be a 2D array."
_check_mesh_is_valid(mesh)
vertices, connectivity = mesh
face_normals = compute_face_normals(vertices, connectivity, face_template)
etoe, etof, face_perm_idx = compute_topology(connectivity, face_template)
return vertices, connectivity, face_normals, etoe, etof, face_perm_idx
def _check_mesh_is_valid(mesh : Tuple[jnp.ndarray, jnp.ndarray]):
assert mesh[0].ndim == 2, f"Mesh vertices must be a 2D array of [n_vertices, n_dim]. Provided shape: {mesh[0].shape}"
assert mesh[1].ndim == 2, f"Mesh connectivity must be a 2D array of [n_elements, n_vertices_per_element]. Provided shape: {mesh[1].shape}"
[docs]
class GeometricDomain(eqx.Module):
'''
Module representing a geometric domain.
These contain precomputed connectivity, face normals, and adjacency information.
'''
vertices : jnp.ndarray
'''Vertices of the domain. Shape is [n_vertices, n_dim]'''
connectivity : jnp.ndarray
'''Connectivity of the domain. Shape is [n_elements, n_vertices_per_element]'''
face_normals : jnp.ndarray
'''Face normals of the domain. Shape is [n_elements, n_faces_per_element, n_dim]'''
etoe : jnp.ndarray
'''Element-to-element adjacency. Shape is [n_elements, n_faces_per_element]. -1 indicates a boundary face (unconnected to another element).'''
etof : jnp.ndarray
'''Element-to-face adjacency. Shape is [n_elements, n_faces_per_element]. Each face has a unique index in the domain.'''
face_perm_idx : jnp.ndarray
'''Face permutation indices. Shape is [n_elements, n_faces_per_element]. Indicates how the face vertices are permuted between adjacent elements.'''
face_template : jnp.ndarray
'''Face template used to define faces of elements. Shape is [n_faces_per_element, n_vertices_per_face].'''
[docs]
@classmethod
def from_mesh(cls, mesh : Tuple[jnp.ndarray, jnp.ndarray], face_template : jnp.ndarray):
'''
Constructs a GeometricDomain from a mesh and a face template (see `jax_sn.tools.mesh_tools._face_template` for an example template).
Parameters
----------
mesh : tuple
A tuple containing vertices and connectivity of the mesh. vertices is a jnp.ndarray of shape [n_vertices, n_dim] and connectivity is a jnp.ndarray of shape [n_elements, n_vertices_per_element].
face_template : jnp.ndarray [n_faces_per_element, n_vertices_per_face]
A 2D array representing the face template.
Returns
-------
domain : GeometricDomain
The constructed GeometricDomain.
'''
return cls(
*_compute_domain_setup_from_mesh(mesh, face_template),
face_template=face_template
)
[docs]
@classmethod
def from_mesh_numpy(cls, mesh : Tuple[jnp.ndarray, jnp.ndarray], face_template : jnp.ndarray):
'''
Constructs a GeometricDomain from a mesh and a face template (see `jax_sn.tools.mesh_tools._face_template` for an example template).
Uses a numpy implementation of the domain setup. This avoids compilation overheads of especially the lexsort, which can take significant time.
Parameters
----------
mesh : tuple
A tuple containing vertices and connectivity of the mesh. vertices is a jnp.ndarray of shape [n_vertices, n_dim] and connectivity is a jnp.ndarray of shape [n_elements, n_vertices_per_element].
face_template : jnp.ndarray [n_faces_per_element, n_vertices_per_face]
A 2D array representing the face template.
Returns
-------
domain : GeometricDomain
The constructed GeometricDomain.
'''
from jax_sn.tools.mesh_tools_numpy import _compute_domain_setup_from_mesh_numpy
domain_setup = _compute_domain_setup_from_mesh_numpy(mesh, face_template)
return cls(
*[jnp.array(x) for x in domain_setup],
face_template=jnp.array(face_template)
)
@property
def n_elements(self):
'''Number of elements in the domain'''
return self.connectivity.shape[0]
@property
def n_faces_per_element(self):
'''Number of faces per element in the domain'''
return self.face_template.shape[0]
[docs]
def containing_cell(self, points : jnp.ndarray):
'''
Containing cell for each point in points.
Uses a cached BVH, so cannot be used in a jitted function
Parameters
----------
points : jnp.ndarray [n_points, n_dim]
The points for which to find the containing cell.
Returns
-------
element_indices : jnp.ndarray [n_points]
The indices of the containing cell for each point. -1 if no containing cell is found
'''
return _containing_cell(self, points)
from .cross_sections import CrossSectionData, CrossSectionDataSingleGroup, CrossSectionsIndexed
[docs]
class Domain(GeometricDomain):
'''
Geometric domain with indexed cross sections.
This is seperated from GeometricDomain because many algorithms do not depend on cross sections and therefore we want to keep the cross sections distinct.
'''
cross_sections_indexed : CrossSectionsIndexed
'''Indexed cross sections for the domain.'''
[docs]
@classmethod
def from_mesh_and_cross_sections(cls, mesh : Tuple[jnp.ndarray, jnp.ndarray], face_template : jnp.ndarray, cross_sections_indexed : CrossSectionsIndexed):
'''
Constructs a Domain from a mesh, face template, and indexed cross sections. Checks compatibility between the mesh and cross sections.
Parameters
----------
mesh : tuple
A tuple containing vertices and connectivity of the mesh. vertices is a jnp.ndarray of shape [n_vertices, n_dim] and connectivity is a jnp.ndarray of shape [n_elements, n_vertices_per_element].
face_template : jnp.ndarray [n_faces_per_element, n_vertices_per_face]
A 2D array representing the face template.
cross_sections_indexed : CrossSectionsIndexed
Indexed cross sections for the domain.
Returns
-------
domain : Domain
The constructed Domain.
'''
_check_cross_section_compatibility(mesh[1].shape[0], cross_sections_indexed)
return cls(
*_compute_domain_setup_from_mesh(mesh, face_template),
face_template = face_template,
cross_sections_indexed = cross_sections_indexed
)
[docs]
@classmethod
def from_mesh_and_cross_sections_numpy(cls, mesh : Tuple[jnp.ndarray, jnp.ndarray], face_template : jnp.ndarray, cross_sections_indexed : CrossSectionsIndexed):
'''
Constructs a Domain from a mesh, face template, and indexed cross sections. Checks compatibility between the mesh and cross sections.
Uses a numpy implementation of the domain setup. This avoids compilation overheads of especially the lexsort, which can take significant time.
Parameters
----------
mesh : tuple
A tuple containing vertices and connectivity of the mesh. vertices is a jnp.ndarray of shape [n_vertices, n_dim] and connectivity is a jnp.ndarray of shape [n_elements, n_vertices_per_element].
face_template : jnp.ndarray [n_faces_per_element, n_vertices_per_face]
A 2D array representing the face template.
cross_sections_indexed : CrossSectionsIndexed
Indexed cross sections for the domain.
Returns
-------
domain : Domain
The constructed Domain.
'''
from jax_sn.tools.mesh_tools_numpy import _compute_domain_setup_from_mesh_numpy
_check_cross_section_compatibility(mesh[1].shape[0], cross_sections_indexed)
domain_setup = _compute_domain_setup_from_mesh_numpy(mesh, face_template)
return cls(
*[jnp.array(x) for x in domain_setup],
face_template=jnp.array(face_template),
cross_sections_indexed=cross_sections_indexed
)
[docs]
@classmethod
def from_geometric_domain_and_cross_sections(cls, geometric_domain : GeometricDomain, cross_sections_indexed : CrossSectionsIndexed):
'''
Constructs a Domain from a GeometricDomain and indexed cross sections. Checks compatibility between the geometric domain and cross sections.
Parameters
----------
geometric_domain : GeometricDomain
The geometric domain.
cross_sections_indexed : CrossSectionsIndexed
Indexed cross sections for the domain.
Returns
-------
domain : Domain
The constructed Domain.
'''
assert isinstance(geometric_domain, GeometricDomain), f"geometric_domain must be an instance of GeometricDomain. Provided type: {type(geometric_domain)}"
_check_cross_section_compatibility(geometric_domain.connectivity.shape[0], cross_sections_indexed)
return cls(geometric_domain.vertices,
geometric_domain.connectivity,
geometric_domain.face_normals,
geometric_domain.etoe,
geometric_domain.etof,
geometric_domain.face_perm_idx,
geometric_domain.face_template,
cross_sections_indexed=cross_sections_indexed
)
def _check_cross_section_compatibility(n_elements : int, cross_sections : CrossSectionsIndexed):
'''
Checks whether the cross sections are compatible with the geometric domain.
Only is static array shape checking, not values.
'''
assert cross_sections.indices.ndim == 1, f"Cross section indices must be a 1D array [ n_elements ]. Provided shape: {cross_sections.indices.shape}"
assert cross_sections.indices.shape[0] == n_elements, f"Cross section indices must have length equal to number of elements in geometric domain. Number of elements in geometric domain: {n_elements}, provided length: {cross_sections.indices.shape[0]}"
cross_sections.cross_section_data.assert_valid_batched()
def _containing_cell(geometric_domain : GeometricDomain, points : jnp.ndarray):
'''
Containing cell for each point in points.
Uses a cached BVH, so cannot be used in a jitted function
Parameters
----------
geometric_domain : GeometricDomain
The geometric domain for which to find the containing cells.
points : jnp.ndarray [n_points, n_dim]
The points for which to find the containing cell.
Returns
-------
element_indices : jnp.ndarray [n_points]
The indices of the containing cell for each point. -1 if no containing cell is found
'''
if isinstance(points, jax.core.Tracer):
raise RuntimeError("The function 'containing_cell' cannot be used inside jitted functions because it relies on a cached BVH.")
from jax_sn.tools.mesh_tools import create_bvh_cached, probe_bvh
bvh = create_bvh_cached(geometric_domain.vertices, geometric_domain.connectivity)
possible_hits = probe_bvh(bvh, points)
elements_containing_point = _check_point_in_element_vec(points[..., None, :], possible_hits, geometric_domain)
element_indices = jnp.where(elements_containing_point, possible_hits, -1)
first_elements = jax.vmap(lambda x : x[jnp.where(x>=0, size=1)[0]], in_axes=0)(element_indices)
return first_elements[:,0]
def _check_point_in_element(point : jnp.ndarray, element : int, geometric_domain : GeometricDomain):
'''
Checks whether a point is inside the element of the geometric domain.
Assumes the domain contains *convex* elements.
Checks whether the point is inside all half-spaces defined by the elements faces.
Parameters
----------
point : jnp.ndarray [n_dim]
The point to check.
element : int
The index of the element to check against.
geometric_domain : GeometricDomain
The geometric domain containing the element.
Returns
-------
inside : bool
True if the point is inside the element, False otherwise.
'''
vertices_ref_face = geometric_domain.connectivity[element, geometric_domain.face_template[:,0]] #[n_faces]
x_min_ref_face = point - geometric_domain.vertices[vertices_ref_face] #[n_faces, n_dim]
face_normals = geometric_domain.face_normals[element] # n_faces, n_dim
inside = jnp.sum(x_min_ref_face * face_normals[:, :x_min_ref_face.shape[-1]], axis=-1) <= 0.0
return jnp.all(inside)
_check_point_in_element_vec = jnp.vectorize(_check_point_in_element, excluded=(2,), signature='(n_dim),()->()')