liblaf.apple.jax.model.dirichlet

Used by:

• Reference Liblaf apple
  •  apple  Dirichlet
  • jax
    •  jax  Dirichlet
    • model
      •  model  Dirichlet
      • dirichlet  dirichlet  Dirichlet

Classes:

• Dirichlet –
• DirichletBuilder –

Dirichlet

Parameters:

• dim

  (int) –

• dirichlet_index

  (Integer[Array, dirichlet]) –

• dirichlet_value

  (Float[Array, dirichlet]) –

• free_index

  (Integer[Array, free]) –

• n_points

  (int) –

Returned by:

• Reference Liblaf apple
  •  apple  Model  dirichlet
  • model  model  Model  dirichlet

Used by:

• Reference Liblaf apple
  •  apple  Model
  • model  model  Model

Methods:

• get_dirichlet –
• get_free –
• set_dirichlet –
• set_free –
• to_full –

Attributes:

• dim (int) –
• dirichlet_index (Integer[Array, ' dirichlet']) –
• dirichlet_value (Float[Array, ' dirichlet']) –
• free_index (Integer[Array, ' free']) –
• n_dirichlet (int) –
• n_free (int) –
• n_full (int) –
• n_points (int) –

dim instance-attribute

dim: int

dirichlet_index instance-attribute

dirichlet_index: Integer[Array, ' dirichlet']

dirichlet_value instance-attribute

dirichlet_value: Float[Array, ' dirichlet']

free_index instance-attribute

free_index: Integer[Array, ' free']

n_dirichlet property

n_dirichlet: int

n_free property

n_free: int

n_full property

n_full: int

n_points instance-attribute

n_points: int

get_dirichlet

get_dirichlet(
    full: Float[Array, "points dim"],
) -> Float[Array, " dirichlet"]

Source code in src/liblaf/apple/jax/model/dirichlet/_dirichlet.py

@eqx.filter_jit
def get_dirichlet(
    self, full: Float[Array, "points dim"]
) -> Float[Array, " dirichlet"]:
    return full.flatten()[self.dirichlet_index]

get_free

get_free(
    full: Float[Array, "points dim"],
) -> Float[Array, " free"]

Source code in src/liblaf/apple/jax/model/dirichlet/_dirichlet.py

@eqx.filter_jit
def get_free(self, full: Float[Array, "points dim"]) -> Float[Array, " free"]:
    return full.flatten()[self.free_index]

set_dirichlet

set_dirichlet(
    full: Float[Array, "points dim"],
    values: Float[ArrayLike, " dirichlet"] | None = None,
) -> Float[Array, "points dim"]

Source code in src/liblaf/apple/jax/model/dirichlet/_dirichlet.py

@eqx.filter_jit
def set_dirichlet(
    self,
    full: Float[Array, "points dim"],
    values: Float[ArrayLike, " dirichlet"] | None = None,
) -> Float[Array, "points dim"]:
    if values is None:
        values = self.dirichlet_value
    return full.flatten().at[self.dirichlet_index].set(values).reshape(full.shape)

set_free

set_free(
    full: Float[Array, "points dim"],
    values: Float[ArrayLike, " free"],
) -> Float[Array, "points dim"]

Source code in src/liblaf/apple/jax/model/dirichlet/_dirichlet.py

@eqx.filter_jit
def set_free(
    self, full: Float[Array, "points dim"], values: Float[ArrayLike, " free"]
) -> Float[Array, "points dim"]:
    return full.flatten().at[self.free_index].set(values).reshape(full.shape)

to_full

to_full(
    free: Float[Array, " free"],
    dirichlet: Float[ArrayLike, " dirichlet"] | None = None,
) -> Float[Array, "points dim"]

Source code in src/liblaf/apple/jax/model/dirichlet/_dirichlet.py

@eqx.filter_jit
def to_full(
    self,
    free: Float[Array, " free"],
    dirichlet: Float[ArrayLike, " dirichlet"] | None = None,
) -> Float[Array, "points dim"]:
    full: Float[Array, "points dim"] = jnp.empty(
        (self.n_points, self.dim), free.dtype
    )
    full = self.set_free(full, free)
    full = self.set_dirichlet(full, dirichlet)
    return full

DirichletBuilder

DirichletBuilder(dim: int = 3)

Parameters:

• mask

  (Bool[ndarray, 'points dim']) –

• value

  (Float[ndarray, 'points dim']) –

Returned by:

• Reference Liblaf apple
  •  apple  ModelBuilder  dirichlet
  • model  model  ModelBuilder  dirichlet

Used by:

• Reference Liblaf apple
  •  apple  ModelBuilder
  • model  model  ModelBuilder  dirichlet

Methods:

• add_pyvista –
• finalize –
• resize –

Attributes:

• dim (int) –
• mask (Bool[ndarray, 'points dim']) –
• n_points (int) –
• value (Float[ndarray, 'points dim']) –

Source code in src/liblaf/apple/jax/model/dirichlet/_builder.py

def __init__(self, dim: int = 3) -> None:
    mask: Bool[np.ndarray, "points dim"] = np.empty((0, dim), bool)
    value: Float[np.ndarray, "points dim"] = np.empty((0, dim))
    self.__attrs_init__(mask=mask, value=value)  # pyright: ignore[reportAttributeAccessIssue]

dim property

dim: int

mask instance-attribute

mask: Bool[ndarray, 'points dim']

n_points property

n_points: int

value instance-attribute

value: Float[ndarray, 'points dim']

add_pyvista

add_pyvista(obj: DataSet) -> None

Source code in src/liblaf/apple/jax/model/dirichlet/_builder.py

def add_pyvista(self, obj: pv.DataSet) -> None:
    point_id = obj.point_data[POINT_ID]
    self.resize(point_id.max() + 1)
    dirichlet_mask: Bool[Array, "points dim"] = self._left_broadcast_to(
        obj.point_data[DIRICHLET_MASK], obj.n_points
    )
    dirichlet_value: Float[Array, "points dim"] = self._left_broadcast_to(
        obj.point_data[DIRICHLET_VALUE], obj.n_points
    )
    self.mask[point_id] = dirichlet_mask
    self.value[point_id] = dirichlet_value

finalize

finalize() -> Dirichlet

Source code in src/liblaf/apple/jax/model/dirichlet/_builder.py

def finalize(self) -> Dirichlet:
    mask: Bool[Array, "points dim"] = jnp.asarray(self.mask)
    dirichlet_index: Integer[Array, " dirichlet"] = jnp.flatnonzero(mask)
    return Dirichlet(
        dim=self.dim,
        dirichlet_index=dirichlet_index,
        dirichlet_value=jnp.asarray(self.value.flat[dirichlet_index]),
        free_index=jnp.flatnonzero(~mask),
        n_points=self.n_points,
    )

resize

resize(n_points: int) -> None

Source code in src/liblaf/apple/jax/model/dirichlet/_builder.py

def resize(self, n_points: int) -> None:
    pad_after: int = n_points - self.n_points
    if pad_after <= 0:
        return
    self.mask = np.pad(self.mask, ((0, pad_after), (0, 0)), constant_values=False)
    self.value = np.pad(self.value, ((0, pad_after), (0, 0)), constant_values=0.0)