aimz 0.1.0__py3-none-any.whl

aimz/__init__.py

@@ -0,0 +1,26 @@
+# Copyright 2025 Eli Lilly and Company
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Initialize package with global setup."""
+
+import logging
+import sys
+
+logger = logging.getLogger(__name__)
+
+if not logger.hasHandlers():
+    logger.setLevel(logging.INFO)
+    handler = logging.StreamHandler(sys.stdout)
+    logger.addHandler(handler)
+    logging.getLogger("py.warnings").addHandler(handler)

aimz/_exceptions.py

@@ -0,0 +1,25 @@
+# Copyright 2025 Eli Lilly and Company
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Custom warnings and errors."""
+
+__all__ = ["KernelValidationError", "NotFittedError"]
+
+
+class NotFittedError(ValueError, AttributeError):
+    """Exception class to raise if model is used before fitting."""
+
+
+class KernelValidationError(Exception):
+    """Exception class to raise if kernel validation fails."""

aimz/data/__init__.py

@@ -0,0 +1,19 @@
+# Copyright 2025 Eli Lilly and Company
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Utilities for data handling and processing."""
+
+from aimz.data.array_loader import ArrayLoader
+
+__all__ = ["ArrayLoader"]

aimz/data/_sharding.py

@@ -0,0 +1,230 @@
+# Copyright 2025 Eli Lilly and Company
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Module for creating functions for sharding.
+
+NOTE: This module is experimental and subject to change. It utilizes JAX's `shard_map()`
+to distribute computations across devices. Tested on CPU and GPU.
+"""
+
+from functools import partial
+from typing import TYPE_CHECKING
+
+from jax import jit
+from jax.experimental.shard_map import shard_map
+from jax.sharding import PartitionSpec
+from numpyro.infer import log_likelihood as log_lik
+
+from aimz.sampling._forward import _sample_forward
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    import jax
+    from jax.sharding import Mesh
+
+
+def _create_sharded_sampler(
+    mesh: "Mesh | None",
+    n_kwargs_array: int,
+    n_kwargs_extra: int,
+) -> "Callable":
+    """Create a sharded posterior predictive sampling function.
+
+    Args:
+        mesh (Mesh): The JAX mesh object defining the device mesh for sharding.
+        n_kwargs_array (int): The number of arguments in the keyword arguments that
+            are array-like (sharded).
+        n_kwargs_extra (int): The number of extra keyword arguments that are not
+            array-like (not sharded).
+
+    Returns:
+        Callable: A sharded function that takes the following arguments:
+            - rng_key (jax.Array): A pseudo-random number generator key.
+            - kernel (Callable): A probabilistic model with Pyro primitives.
+            - posterior_samples (dict): A dictionary of posterior samples.
+            - batch_shape (tuple[int]): The shape of the batch dimension, specifically
+                `(num_samples,)`.
+            - param_input (str): The name of the parameter in the `kernel` for the
+                input data.
+            - kwargs_key (tuple[str]): A tuple of keyword argument names.
+            - X (jax.Array): Input data.
+            - *args (tuple): Additional arguments constructed from the original keyword
+                arguments (both sharded and non-sharded).
+
+    """
+
+    def f(
+        kernel: "Callable",
+        num_samples: int,
+        rng_key: "jax.Array",
+        return_sites: tuple[str],
+        posterior_samples: dict[str, "jax.Array"],
+        param_input: str,
+        kwargs_key: tuple[str],
+        X: "jax.Array",
+        *args: tuple,
+    ) -> dict[str, "jax.Array"]:
+        return _sample_forward(
+            model=kernel,
+            num_samples=num_samples,
+            rng_key=rng_key,
+            return_sites=return_sites,
+            posterior_samples=posterior_samples,
+            model_kwargs={
+                param_input: X,
+                **dict(zip(kwargs_key, args, strict=True)),
+            },
+        )
+
+    if mesh is None:
+        return partial(
+            jit,
+            static_argnames=[
+                "kernel",
+                "num_samples",
+                "return_sites",
+                "param_input",
+                "kwargs_key",
+            ],
+        )(f)
+
+    (axis,) = mesh.axis_names
+
+    return partial(
+        jit,
+        static_argnames=[
+            "kernel",
+            "num_samples",
+            "return_sites",
+            "param_input",
+            "kwargs_key",
+        ],
+    )(
+        partial(
+            shard_map,
+            mesh=mesh,
+            in_specs=(
+                None,  # kernel
+                None,  # posterior_samples
+                None,  # rng_key
+                None,  # num_samples
+                None,  # return_sites
+                None,  # param_input
+                None,  # kwargs_key
+                PartitionSpec(axis),  # X
+                *(
+                    [PartitionSpec(axis)] * n_kwargs_array  # kwargs_array
+                    + [None] * n_kwargs_extra  # kwargs_extra
+                ),
+            ),
+            out_specs=PartitionSpec(None, axis),
+            check_rep=False,
+        )(f),
+    )
+
+
+def _create_sharded_log_likelihood(
+    mesh: "Mesh | None",
+    n_kwargs_array: int,
+    n_kwargs_extra: int,
+) -> "Callable":
+    """Create a sharded log-likelihood function.
+
+    Args:
+        mesh (Mesh): The JAX mesh object defining the device mesh for sharding.
+        n_kwargs_array (int): The number of arguments in the keyword arguments that are
+            array-like (sharded).
+        n_kwargs_extra (int): The number of extra keyword arguments that are not
+            array-like (not sharded).
+
+    Returns:
+        Callable: A sharded function that takes the following arguments:
+            - kernel (Callable): A probabilistic model with Pyro primitives optimized
+                with variational inference.
+            - posterior_samples (dict): A dictionary of posterior samples.
+            - param_input (str): The name of the parameter in the `kernel` for the input
+                data.
+            - param_output (str): The name of the parameter in the `kernel` for the
+                output data.
+            - kwargs_key (tuple[str]): A tuple of keyword argument names.
+            - X (jax.Array): Input data.
+            - y (jax.Array): Output data.
+            - *args (tuple): Additional arguments constructed from the original keyword
+                arguments (both sharded and non-sharded).
+
+    """
+
+    def f(
+        kernel: "Callable",
+        posterior_samples: dict,
+        param_input: str,
+        param_output: str,
+        kwargs_key: tuple[str],
+        X: "jax.Array",
+        y: "jax.Array",
+        *args: tuple,
+    ) -> "jax.Array":
+        return log_lik(
+            kernel,
+            posterior_samples=posterior_samples,
+            **{
+                param_input: X,
+                param_output: y,
+                **dict(zip(kwargs_key, args, strict=True)),
+            },
+        ).get(param_output)
+
+    if mesh is None:
+        return partial(
+            jit,
+            static_argnames=[
+                "kernel",
+                "param_input",
+                "param_output",
+                "kwargs_key",
+            ],
+        )(f)
+
+    (axis,) = mesh.axis_names
+
+    return partial(
+        jit,
+        static_argnames=[
+            "kernel",
+            "param_input",
+            "param_output",
+            "kwargs_key",
+        ],
+    )(
+        partial(
+            shard_map,
+            mesh=mesh,
+            in_specs=(
+                None,  # kernel
+                None,  # posterior_samples
+                None,  # param_input
+                None,  # param_output
+                None,  # kwargs_key
+                PartitionSpec(axis),  # X
+                PartitionSpec(axis),  # y
+                *(
+                    [PartitionSpec(axis)] * n_kwargs_array  # kwargs_array
+                    + [None] * n_kwargs_extra  # kwargs_extra
+                ),
+            ),
+            out_specs=PartitionSpec(None, axis),
+            check_rep=False,
+        )(f),
+    )

aimz/data/array_loader.py

@@ -0,0 +1,205 @@
+# Copyright 2025 Eli Lilly and Company
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Module for custom data loader with padding logic for JAX arrays.
+
+This module defines a custom `ArrayLoader` that processes batches of data and applies
+padding to ensure the batch size is compatible with sharding across multiple XLA
+devices.
+"""
+
+from typing import TYPE_CHECKING
+
+import jax.numpy as jnp
+from jax import device_put
+from jax_dataloader import ArrayDataset
+from torch.utils.data import DataLoader
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    import jax
+    from jax.sharding import NamedSharding
+    from torch.utils.data import Sampler
+
+
+class ArrayLoader(DataLoader):
+    """A custom DataLoader class that extends PyTorch's DataLoader."""
+
+    def __init__(
+        self,
+        dataset: ArrayDataset,
+        *,
+        batch_size: int = 1,
+        sampler: "Sampler | None" = None,
+        num_workers: int = 0,
+        collate_fn: "Callable | None" = None,
+        pin_memory: bool = False,
+    ) -> None:
+        """Initializes an ArrayLoader instance."""
+        super().__init__(
+            dataset,
+            batch_size=batch_size,
+            sampler=sampler,
+            num_workers=num_workers,
+            collate_fn=collate_fn,
+            pin_memory=pin_memory,
+        )
+
+    @staticmethod
+    def calculate_padding(batch_size: int, num_devices: int) -> int:
+        """Calculate the number of padding needed.
+
+        Args:
+            batch_size (int): The size of the batch.
+            num_devices (int): The number of devices.
+
+        Returns:
+            int: The number of padding rows (or elements) needed to make the batch size
+                divisible by the number of devices.
+        """
+        remainder = batch_size % num_devices
+        return 0 if remainder == 0 else num_devices - remainder
+
+    @staticmethod
+    def pad_array(x: "jax.Array", n_pad: int, axis: int) -> "jax.Array":
+        """Pad an array to ensure compatibility with sharding.
+
+        Args:
+            x (jax.Array): The input array to be padded.
+            n_pad (int): The number of padding elements to add.
+            axis (int): The axis along which to apply the padding.
+
+        Returns:
+            jax.Array: The padded array with padding applied along the specified axis.
+
+        Raises:
+            ValueError: If padding is requested along an unsupported axis for a 1D
+                array.
+        """
+        if x.ndim == 1:
+            if axis == 0:
+                return jnp.pad(x, pad_width=(0, n_pad), mode="edge")
+            msg = "Padding 1D arrays is only supported along axis 0."
+            raise ValueError(msg)
+
+        # Initialize all axes with no padding
+        pad_width: list[tuple[int, int]] = [(0, 0)] * x.ndim
+        # Apply padding to the specified axis
+        pad_width[axis] = (0, n_pad)
+
+        return jnp.pad(x, pad_width=pad_width, mode="edge")
+
+    @staticmethod
+    def collate_without_output(
+        batch: list[tuple],
+        device: "NamedSharding | None" = None,
+    ) -> tuple:
+        """Collate function to process batches with sharding and padding.
+
+        This function unpacks the batch of data, converts it into JAX arrays, and
+        applies padding to ensure the batch size is compatible with the number of
+        devices, if sharding is necessary. When a device is provided, the data is
+        automatically distributed across the available devices.
+
+        Args:
+            batch (list[tuple]): A list of tuples, where each tuple contains the input
+                data, optional target data, and array-like keyword arguments.
+            device (NamedSharding | None, optional): Sharding using named axes for
+                parallel data distribution across devices. Defaults to `None`, meaning
+                no sharding is applied.
+
+        Returns:
+            tuple: A tuple containing:
+                - n_pad (int): The number of padding rows/elements added (0 if no
+                    padding was required).
+                - x_batch (jax.Array): The input batch with padding applied if
+                    necessary.
+                - kwargs_batch (list[jax.Array]): A list of keyword arguments with
+                    padding applied if necessary.
+        """
+        x_batch, *kwargs_batch = map(jnp.asarray, zip(*batch, strict=True))
+
+        n_pad = (
+            ArrayLoader.calculate_padding(
+                len(x_batch),
+                num_devices=device.num_devices,
+            )
+            if device
+            else 0
+        )
+        if n_pad:
+            x_batch = ArrayLoader.pad_array(x_batch, n_pad=n_pad, axis=0)
+            kwargs_batch = [
+                ArrayLoader.pad_array(x, n_pad=n_pad, axis=0) for x in kwargs_batch
+            ]
+
+        if device:
+            x_batch = device_put(x_batch, device=device)
+            kwargs_batch = [device_put(x, device=device) for x in kwargs_batch]
+
+        return n_pad, x_batch, *kwargs_batch
+
+    @staticmethod
+    def collate_with_sharding(
+        batch: list[tuple],
+        device: "NamedSharding | None" = None,
+    ) -> tuple:
+        """Collate function to process batches with sharding and padding.
+
+        This function unpacks the batch of data, converts it into JAX arrays, and
+        applies padding to ensure the batch size is compatible with the number of
+        devices, if sharding is necessary. When a device is provided, the data is
+        automatically distributed across the available devices.
+
+        Args:
+            batch (list[tuple]): A list of tuples, where each tuple contains the input
+                data, optional target data, and array-like keyword arguments.
+            device (NamedSharding | None, optional): Sharding using named axes for
+                parallel data distribution across devices. Defaults to `None`, meaning
+                no sharding is applied.
+
+        Returns:
+            tuple: A tuple containing:
+                - n_pad (int): The number of padding rows/elements added (0 if no
+                    padding was required).
+                - x_batch (jax.Array): The input batch with padding applied if
+                    necessary.
+                - y_batch (jax.Array): The target batch with padding applied.
+                - kwargs_batch (list[jax.Array]): A list of keyword arguments with
+                    padding applied if necessary.
+        """
+        x_batch, y_batch, *kwargs_batch = map(jnp.asarray, zip(*batch, strict=True))
+
+        n_pad = (
+            ArrayLoader.calculate_padding(
+                len(x_batch),
+                num_devices=device.num_devices,
+            )
+            if device
+            else 0
+        )
+        if n_pad:
+            x_batch = ArrayLoader.pad_array(x_batch, n_pad=n_pad, axis=0)
+            y_batch = ArrayLoader.pad_array(y_batch, n_pad=n_pad, axis=0)
+            kwargs_batch = [
+                ArrayLoader.pad_array(x, n_pad=n_pad, axis=0) for x in kwargs_batch
+            ]
+
+        if device:
+            x_batch = device_put(x_batch, device=device)
+            y_batch = device_put(y_batch, device=device)
+            kwargs_batch = [device_put(x, device=device) for x in kwargs_batch]
+
+        return n_pad, x_batch, y_batch, *kwargs_batch

aimz/model/__init__.py

@@ -0,0 +1,19 @@
+# Copyright 2025 Eli Lilly and Company
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Model object."""
+
+from aimz.model.impact_model import ImpactModel
+
+__all__ = ["ImpactModel"]

aimz/model/_core.py

@@ -0,0 +1,91 @@
+# Copyright 2025 Eli Lilly and Company
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Base class for impact model."""
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Self
+
+from aimz.utils._validation import _validate_kernel_signature
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    import jax
+
+
+class BaseModel(ABC):
+    """Abstract base class for the impact model.
+
+    Attributes:
+        kernel (Callable): A probabilistic model with Pyro primitives.
+    """
+
+    def __init__(
+        self,
+        kernel: "Callable",
+        param_input: str = "X",
+        param_output: str = "y",
+    ) -> None:
+        """Initialize the BaseModel with a callable model.
+
+        Args:
+            kernel (Callable): A probabilistic model with Pyro primitives.
+            param_input (str, optional): Name of the parameter in the `kernel` for the
+                main input data. Defaults to `"X"`.
+            param_output (str, optional): Name of the parameter in the `kernel` for the
+                output data. Defaults to `"y"`.
+        """
+        self.kernel = kernel
+        self.param_input = param_input
+        self.param_output = param_output
+        _validate_kernel_signature(self.kernel, self.param_input, self.param_output)
+
+    @abstractmethod
+    def fit(
+        self,
+        X: "jax.Array",
+        y: "jax.Array",
+        **kwargs: object,
+    ) -> Self:
+        """Fit the model to the input data `X` and output data `y`.
+
+        Args:
+            X (jax.Array): Input data with shape `(n_samples_X, n_features)`.
+            y (jax.Array): Output data with shape `(n_samples_Y,)`.
+            **kwargs (object): Additional arguments passed to the model, except for `X`
+                and `y`. All array-like objects in `**kwargs` are expected to be JAX
+                arrays.
+
+        Returns:
+            BaseModel: The fitted model instance, enabling method chaining.
+        """
+        return self
+
+    @abstractmethod
+    def predict(self, X: "jax.Array", **kwargs: object) -> None:
+        """Predict the output based on the fitted model.
+
+        Args:
+            X (jax.Array): Input data with shape `(n_samples_X, n_features)`.
+            **kwargs (object): Additional arguments passed to the model, except for `X`
+                and `y`. All array-like objects in `**kwargs` are expected to be JAX
+                arrays.
+        """
+
+    @abstractmethod
+    def estimate_effect(
+        self,
+    ) -> None:
+        """Estimate the effect of an intervention."""