Fix TFP Version Incompatibility with Transformed Distributions (#64)

* Replaced SquareRoot bijector Transformation with builtin `tfb.Power`

* Remove py36 support

.github/workflows/run-CI.yml

@@ -6,16 +6,12 @@ jobs:
     strategy:
       matrix:
         os: [ubuntu-latest, macos-latest, windows-latest]
-        python: [3.6, 3.7, 3.8, 3.9]
+        python: [3.7, 3.8, 3.9]
         exclude:
-          - os: macos-latest
-            python: 3.6
           - os: macos-latest
             python: 3.7
           - os: macos-latest
             python: 3.8
-          - os: windows-latest
-            python: 3.6
           - os: windows-latest
             python: 3.7
           - os: windows-latest

causalimpact/__version__.py

@@ -13,4 +13,4 @@
 # limitations under the License.
 
 
-__version__ = '0.0.11'
+__version__ = '0.0.12-rc.0'

causalimpact/model.py

@@ -226,8 +226,8 @@ def build_bijector(dist: tfd.Distribution) -> tfd.Distribution:
       new_dist: tfd.Distribution
           New distribution given by `y = G(X)`.
     """
-    bijector = SquareRootBijector()
-    new_dist = tfd.TransformedDistribution(dist, bijector)
+    sqrt_bi = tfb.Power(.5)
+    new_dist = tfd.TransformedDistribution(dist, sqrt_bi)
     return new_dist
 
 
@@ -443,104 +443,3 @@ def build_posterior_dist(
         parameter_samples=parameter_samples,
         num_steps_forecast=num_steps_forecast
     )
-
-
-class SquareRootBijector(tfb.Bijector):
-    """
-    Compute `Y = g(X) = X ** (1 / 2) which transforms variance into standard deviation.
-    Main reference for building this bijector is the original [PowerTransform](https://github.com/tensorflow/probability/blob/v0.11.1/tensorflow_probability/python/bijectors/power_transform.py) # noqa: E501
-    """
-    def __init__(
-        self,
-        validate_args: bool = False,
-        name: str = 'square_root_bijector'
-    ):
-        """
-        Args
-        ----
-          validate_args: bool
-              Indicates whether arguments should be checked for correctness.
-          name: str
-              Name given to ops managed by this object.
-        """
-        # Without these `parameters` the code won't be compatible with future versions
-        # of tfp:
-        # https://github.com/tensorflow/probability/issues/1202
-        parameters = dict(locals())
-        with tf.name_scope(name) as name:
-            super().__init__(
-                forward_min_event_ndims=0,
-                validate_args=validate_args,
-                parameters=parameters,
-                name=name)
-
-    def _forward(self, x: Union[float, np.array, tf.Tensor]) -> tf.Tensor:
-        """
-        Implements the forward pass `G` as given by `Y = G(X)`. In this case, it's a
-        simple square root of X.
-
-        Args
-        ----
-          x: Union[float, np.array, tf.Tensor])
-              Variable `X` to receive the transformation.
-
-        Returns
-        -------
-          X: tf.Tensor
-              Square root of `x`.
-        """
-        return tf.sqrt(x)
-
-    def _inverse(self, y: Union[float, np.array, tf.Tensor]) -> tf.Tensor:
-        """
-        Implements G^-1(y).
-
-        Args
-        ----
-          y: Union[float, np.array, tf.Tensor]
-              Values to be transformed back. In this case, they will be squared.
-
-        Returns
-        -------
-          y: tf.Tensor
-              Squared `y`.
-        """
-        return tf.square(y)
-
-    def _inverse_log_det_jacobian(self, y: tf.Tensor) -> tf.Tensor:
-        """
-        When transforming from `P(X)` to `P(Y)` it's necessary to compute the log of the
-        determinant of the Jacobian matrix for each correspondent function `G` which
-        accounts for the volumetric transformations on each domain.
-
-        The inverse log determinant is given by:
-
-        `ln(|J(G^-1(Y)|) = ln(|J(Y ** 2)|) = ln(|2 * Y|) = ln(2 * Y)`
-
-        Args
-        ----
-          y: tf.Tensor
-
-        Returns
-        -------
-          tf.Tensor
-        """
-        return tf.math.log(2 * y)
-
-    def _forward_log_det_jacobian(self, x: tf.Tensor) -> tf.Tensor:
-        """
-        Computes the volumetric change when moving forward from `P(X)` to `P(Y)`, given
-        by:
-
-        `ln(|J(G(X))|) = ln(|J(sqrt(X))|) = ln(|(1 / 2) * X ** (-1 / 2)|) =
-                       = (-1 / 2) * ln(4.0 * X)
-
-        Args
-        ----
-          x: tf.Tensor
-
-        Returns
-        -------
-          tf.tensor
-        """
-        return -0.5 * tf.math.log(4.0 * x)

setup.py

@@ -47,9 +47,9 @@
 
 install_requires = [
     'jinja2',
-    'pandas <= 1.3.5',
-    'tensorflow',
-    'tensorflow-probability == 0.14.0',
+    'pandas >= 1.3.5',
+    'tensorflow >= 2.10',
+    'tensorflow-probability >= 0.18',
     'matplotlib',
 ]
 tests_require = [

tests/test_main.py

@@ -31,6 +31,10 @@
 from causalimpact import CausalImpact
 from causalimpact.misc import standardize
 
+seed = 123
+tf.random.set_seed(seed)
+np.random.seed(seed)
+
 
 def test_default_causal_cto(rand_data, pre_int_period, post_int_period):
     ci = CausalImpact(rand_data, pre_int_period, post_int_period)

tests/test_model.py

@@ -153,23 +153,6 @@ def test_check_input_model():
         cimodel.check_input_model(model, pre_data, post_data)
 
 
-def test_SquareRootBijector():
-    bijector = cimodel.SquareRootBijector()
-    assert bijector.name == 'square_root_bijector'
-    x = np.array([3.0, 4.0])
-    y = np.array([2.0, 3.0])
-    np.testing.assert_almost_equal(bijector.forward(x), np.sqrt(x))
-    np.testing.assert_almost_equal(bijector.inverse(y), np.square(y))
-    np.testing.assert_almost_equal(
-        bijector.forward_log_det_jacobian(x, event_ndims=0),
-        -.5 * np.log(4.0 * x)
-    )
-    np.testing.assert_almost_equal(
-        bijector.inverse_log_det_jacobian(y, event_ndims=0),
-        np.log(2 * y)
-    )
-
-
 def test_build_default_model(rand_data, pre_int_period, post_int_period):
     prior_level_sd = 0.01
     pre_data = pd.DataFrame(rand_data.iloc[pre_int_period[0]: pre_int_period[1], 0])
@@ -180,11 +163,11 @@ def test_build_default_model(rand_data, pre_int_period, post_int_period):
     assert isinstance(model, tfp.sts.Sum)
     obs_prior = model.parameters[0].prior
     assert isinstance(obs_prior, tfd.TransformedDistribution)
-    assert isinstance(obs_prior.bijector, cimodel.SquareRootBijector)
+    assert isinstance(obs_prior.bijector, tfp.bijectors.Power)
     assert isinstance(obs_prior.distribution, tfd.InverseGamma)
     level_prior = model.parameters[1].prior
     assert isinstance(level_prior, tfd.TransformedDistribution)
-    assert isinstance(level_prior.bijector, cimodel.SquareRootBijector)
+    assert isinstance(level_prior.bijector, tfp.bijectors.Power)
     assert isinstance(level_prior.distribution, tfd.InverseGamma)
     assert level_prior.dtype == tf.float32
 
@@ -196,11 +179,11 @@ def test_build_default_model(rand_data, pre_int_period, post_int_period):
     assert isinstance(model, tfp.sts.Sum)
     obs_prior = model.parameters[0].prior
     assert isinstance(obs_prior, tfd.TransformedDistribution)
-    assert isinstance(obs_prior.bijector, cimodel.SquareRootBijector)
+    assert isinstance(obs_prior.bijector, tfp.bijectors.Power)
     assert isinstance(obs_prior.distribution, tfd.InverseGamma)
     level_prior = model.parameters[1].prior
     assert isinstance(level_prior, tfd.TransformedDistribution)
-    assert isinstance(level_prior.bijector, cimodel.SquareRootBijector)
+    assert isinstance(level_prior.bijector, tfp.bijectors.Power)
     assert isinstance(level_prior.distribution, tfd.InverseGamma)
     assert level_prior.dtype == tf.float32
     linear = model.components[1]
@@ -218,12 +201,12 @@ def test_build_default_model(rand_data, pre_int_period, post_int_period):
     assert isinstance(model, tfp.sts.Sum)
     obs_prior = model.parameters[0].prior
     assert isinstance(obs_prior, tfd.TransformedDistribution)
-    assert isinstance(obs_prior.bijector, cimodel.SquareRootBijector)
+    assert isinstance(obs_prior.bijector, tfp.bijectors.Power)
     assert isinstance(obs_prior.distribution, tfd.InverseGamma)
     assert obs_prior.dtype == tf.float32
     level_prior = model.parameters[1].prior
     assert isinstance(level_prior, tfd.TransformedDistribution)
-    assert isinstance(level_prior.bijector, cimodel.SquareRootBijector)
+    assert isinstance(level_prior.bijector, tfp.bijectors.Power)
     assert isinstance(level_prior.distribution, tfd.InverseGamma)
     assert level_prior.dtype == tf.float32
     linear = model.components[1]

tox.ini

@@ -1,11 +1,10 @@
 [tox]
 envlist =
-    py{36, 37, 38, 39}-{linux,macos,windows}
+    py{37, 38, 39}-{linux,macos,windows}
     gh-actions-coveralls
 
 [gh-actions]
 python =
-    3.6: py36
     3.7: py37
     3.8: py38
     3.9: py39, lint, isort-check, coverage
@@ -26,7 +25,7 @@ deps =
     -rtest-requirements.txt
 commands =
     # To run specific test, simply run `tox -e py39 -- tests/test_data.py -k test_input_data`
-    python -m pytest {posargs: tests/} -m "not slow" -n 2 -x
+    python -m pytest {posargs: tests/} -m "not slow" -n 4 -x
 
 [testenv:isort]
 deps =
@@ -53,13 +52,13 @@ deps =
     {[testenv]deps}
     pytest-cov
 commands =
-    python -m pytest tests --cov-report html --cov causalimpact -n 2 -x
+    python -m pytest tests --cov-report html --cov causalimpact -n 4 -x
 
 [testenv:GHA-coverage]
 deps =
     {[testenv]deps}
     pytest-cov
     coverage
 commands =
-    python -m pytest tests --cov-report xml --cov causalimpact -n 2 -x
+    python -m pytest tests --cov-report xml --cov causalimpact -n 4 -x
     coverage lcov