Updated regions tests.

mccube/regions.py

@@ -1,4 +1,4 @@
-"""Defines abstract and conrete integration regions."""
+"""Defines abstract and concrete integration regions."""
 from __future__ import annotations
 
 import abc
@@ -13,12 +13,9 @@ class AbstractIntegrationRegion(eqx.Module):
 
  Attributes:
  dimension: dimension $d$ of the integration region $\Omega$.
- affine_transformation_matrix: a matrix specifying an affine transformation of
- the integration region.
  """
 
  dimension: int
- affine_transformation_matrix: Float[ArrayLike, "d+1 d+1"]
 
  @abc.abstractmethod
  def weight_function(x: ArrayLike) -> Array:
@@ -30,6 +27,11 @@ def volume(self) -> float:
  """Volume of the weighted integration region."""
  ...
 
+ @abc.abstractproperty
+ def affine_transformation_matrix(self) -> Float[ArrayLike, "d+1 d+1"]:
+ """Affine transformation matrix from the canonical region."""
+ ...
+
 
 class GaussianIntegrationRegion(AbstractIntegrationRegion):
  r"""d-dimensional unnormalized gaussian weighted Euclidean integration region.
@@ -39,9 +41,13 @@ class GaussianIntegrationRegion(AbstractIntegrationRegion):
 
  Attributes:
  dimension: dimension $d$ of the integration region $\Omega$.
+ mean: mean parameter for the Gaussian weight function.
+ covariance: covariance parameter for the Gaussian weight function.
  affine_transformation_matrix: a matrix specifying an affine transformation of
  the integration region.
  """
+ mean: Float[ArrayLike, " d"]
+ covariance: Float[ArrayLike, "d d"]
 
  def __init__(
  self,
@@ -58,9 +64,8 @@ def __init__(
  implicitly diag(1/2) if None.
  """
  self.dimension = dimension
- self.affine_transformation_matrix = self.construct_affine_transformation_matrix(
- dimension, mean, covariance
- )
+ self.mean = np.ones(dimension) if mean is None else mean
+ self.covariance = np.eye(dimension) / 2 if covariance is None else covariance
 
  def weight_function(self, x: ArrayLike) -> Array:
  return np.exp(-np.sum(x**2, axis=-1))
@@ -69,20 +74,14 @@ def weight_function(self, x: ArrayLike) -> Array:
  def volume(self) -> float:
  return np.pi ** (self.dimension / 2)
 
- @staticmethod
- def construct_affine_transformation_matrix(
- dimension: int,
- mean: None | Float[ArrayLike, " d"] = None,
- covariance: None | Float[ArrayLike, "d d"] = None,
- ):
- """Generate affine transform for desired Gaussian mean and covariance."""
- default_cov = np.eye(dimension) / 2
- target_cov = covariance if covariance is not None else default_cov
-
- default_mean = np.zeros(dimension)
- target_mean = mean if mean is not None else default_mean
+ @property
+ def affine_transformation_matrix(self):
+ default_cov = np.eye(self.dimension) / 2
+ target_cov = self.covariance
 
- default_transform = np.eye(dimension + 1)
+ default_mean = np.zeros(self.dimension)
+ target_mean = self.mean
+ default_transform = np.eye(self.dimension + 1)
  default_transform[1:, 0] = default_mean
  default_transform[1:, 1:] = default_cov
 

tests/test_regions.py

@@ -0,0 +1,73 @@
+import chex
+
+import numpy as np
+
+from absl.testing import absltest, parameterized
+
+from mccube.regions import _psd_quadratic_transformation, GaussianIntegrationRegion
+
+
+class RegionTest(chex.TestCase):
+ """
+ Ensure that cubature regions can be instantiated correctly, have appropriate
+ attributes set, and correctly define their affine transformations.
+ """
+
+ # fmt: off
+ @parameterized.named_parameters(
+ ("1D (Scalar)" , 1, 3 * np.ones(1), 4 * np.eye(1)),
+ ("2D (None Mean)" , 2, None, 2 * np.eye(2)),
+ ("3D (None Cov)" , 3, 2 * np.ones(3), None),
+ ("4D (None Mean + Cov)", 3, 2 * np.ones(3), None),
+ )
+ # fmt: on
+ def test_gaussian(self, dimension, mean, covariance):
+ region = GaussianIntegrationRegion(dimension, mean, covariance)
+ self.assertEqual(region.dimension, dimension)
+ test_mean = np.ones(dimension) if mean is None else mean
+ test_covariance = np.eye(dimension) / 2 if covariance is None else covariance
+ chex.assert_trees_all_equal(region.mean, test_mean)
+ chex.assert_trees_all_equal(region.covariance, test_covariance)
+
+ test_affine = np.eye(dimension + 1)
+ test_affine[1:, 1:] = test_covariance
+ test_affine[1:, 0] = test_mean
+
+ canonical_affine = np.eye(dimension + 1)
+ canonical_affine[1:, 1:] = np.eye(dimension) / 2
+ test_transform = _psd_quadratic_transformation(canonical_affine, test_affine)
+ chex.assert_trees_all_equal(region.affine_transformation_matrix, test_transform)
+
+ def test_psd_quadratic_transformation(self):
+ """As per example 2 pg 10-12 :cite:p:`stroud1971`, except $T$ is denoted $M$."""
+ # fmt: off
+ # X as per eq 1.4-10.
+ x = np.array([[1, 1, 0],
+ [1,-1, 0],
+ [1, 0, 1],
+ [1, 0,-1]])
+ A = np.array([[1,0,0],
+ [0,1,0],
+ [0,0,1]])
+ # B is as per eq 1.4-11.
+ B = np.array([[ 50, 0, 0],
+ [-20, 20,-16],
+ [-20,-16, 20]]) / 9
+ M_inv = np.array([[ 1, 0, 0],
+ [-5/3, 4/3,-2/3],
+ [-5/3,-2/3, 4/3]])
+ # fmt: on
+ # Check linear
+ M_linear = _psd_quadratic_transformation(A[1:, 1:], B[1:, 1:], affine=False)
+ chex.assert_tree_all_close(B[1:, 1:], M_linear.T @ A[1:, 1:] @ M_linear)
+ chex.assert_trees_all_close(M_inv[1:, 1:], M_linear)
+ chex.assert_trees_all_close(M_inv[1:, 1:] @ x[:, 1:].T, M_linear @ x[:, 1:].T)
+
+ # Check affine
+ M_affine = _psd_quadratic_transformation(A, B)
+ chex.assert_trees_all_close(M_inv, M_affine)
+ chex.assert_trees_all_close(M_inv @ x.T, M_affine @ x.T)
+
+
+if __name__ == "__main__":
+ absltest.main()