Skip to content

Base Initial Condition

exponax.ic.BaseRandomICGenerator ¤

Bases: Module

Source code in exponax/ic/_base_ic.py 
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
class BaseRandomICGenerator(eqx.Module):
    num_spatial_dims: int
    indexing: str = "ij"

    def gen_ic_fun(self, *, key: PRNGKeyArray) -> BaseIC:
        """
        Generate an initial condition function.

        **Arguments**:
            - `key`: A jax random key.

        **Returns**:
            - `ic`: An initial condition function that can be evaluated at
                degree of freedom locations.
        """
        raise NotImplementedError(
            "This random ic generator cannot represent its initial condition as a function. Directly evaluate it."
        )

    def __call__(
        self,
        num_points: int,
        *,
        key: PRNGKeyArray,
    ) -> Float[Array, "1 ... N"]:
        """
        Generate a random initial condition.

        **Arguments**:
            - `num_points`: The number of grid points in each dimension.
            - `key`: A jax random key.
            - `indexing`: The indexing convention for the grid.

        **Returns**:
            - `u`: The initial condition evaluated at the grid points.
        """
        ic_fun = self.gen_ic_fun(key=key)
        grid = make_grid(
            self.num_spatial_dims,
            self.domain_extent,
            num_points,
            indexing=self.indexing,
        )
        return ic_fun(grid)
num_spatial_dims instance-attribute ¤
num_spatial_dims: int
indexing class-attribute instance-attribute ¤
indexing: str = 'ij'
gen_ic_fun ¤
gen_ic_fun(*, key: PRNGKeyArray) -> BaseIC

Generate an initial condition function.

Arguments: - key: A jax random key.

Returns: - ic: An initial condition function that can be evaluated at degree of freedom locations.

Source code in exponax/ic/_base_ic.py 
28
29
30
31
32
33
34
35
36
37
38
39
40
41
def gen_ic_fun(self, *, key: PRNGKeyArray) -> BaseIC:
    """
    Generate an initial condition function.

    **Arguments**:
        - `key`: A jax random key.

    **Returns**:
        - `ic`: An initial condition function that can be evaluated at
            degree of freedom locations.
    """
    raise NotImplementedError(
        "This random ic generator cannot represent its initial condition as a function. Directly evaluate it."
    )
__call__ ¤
__call__(
    num_points: int, *, key: PRNGKeyArray
) -> Float[Array, "1 ... N"]

Generate a random initial condition.

Arguments: - num_points: The number of grid points in each dimension. - key: A jax random key. - indexing: The indexing convention for the grid.

Returns: - u: The initial condition evaluated at the grid points.

Source code in exponax/ic/_base_ic.py 
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
def __call__(
    self,
    num_points: int,
    *,
    key: PRNGKeyArray,
) -> Float[Array, "1 ... N"]:
    """
    Generate a random initial condition.

    **Arguments**:
        - `num_points`: The number of grid points in each dimension.
        - `key`: A jax random key.
        - `indexing`: The indexing convention for the grid.

    **Returns**:
        - `u`: The initial condition evaluated at the grid points.
    """
    ic_fun = self.gen_ic_fun(key=key)
    grid = make_grid(
        self.num_spatial_dims,
        self.domain_extent,
        num_points,
        indexing=self.indexing,
    )
    return ic_fun(grid)

exponax.ic.BaseIC ¤

Bases: Module, ABC

Source code in exponax/ic/_base_ic.py 
 9
10
11
12
13
14
15
16
17
18
19
20
21
class BaseIC(eqx.Module, ABC):
    @abstractmethod
    def __call__(self, x: Float[Array, "D ... N"]) -> Float[Array, "1 ... N"]:
        """
        Evaluate the initial condition.

        **Arguments**:
            - `x`: The grid points.

        **Returns**:
            - `u`: The initial condition evaluated at the grid points.
        """
        pass
__call__ abstractmethod ¤
__call__(
    x: Float[Array, "D ... N"]
) -> Float[Array, "1 ... N"]

Evaluate the initial condition.

Arguments: - x: The grid points.

Returns: - u: The initial condition evaluated at the grid points.

Source code in exponax/ic/_base_ic.py 
10
11
12
13
14
15
16
17
18
19
20
21
@abstractmethod
def __call__(self, x: Float[Array, "D ... N"]) -> Float[Array, "1 ... N"]:
    """
    Evaluate the initial condition.

    **Arguments**:
        - `x`: The grid points.

    **Returns**:
        - `u`: The initial condition evaluated at the grid points.
    """
    pass