# Path: epi_forecast_stat_mech/statistical_models/linear_model.py
# class LinearParameters(collections.namedtuple(
#     'LinearModelParameters',
#     ('alpha',  # An array of linear weights.
#      'intercept',  # The intercept of a linear model.
#      ))):
# class LinearModel:
#   def init(cls, covariate_dim, prediction_dim, intercept=None):
#   def log_prior(self, parameters):
#   def log_likelihood(self, parameters, covariates, mechanistic_parameters):
#   def predict(self, parameters, covariates):
# 
# Path: epi_forecast_stat_mech/statistical_models/tree_util.py
# def tree_flatten(pytree):
# def tree_unflatten(treedef, leaves):
# def tree_map(f, structure):
# def tree_multimap(f, structure, *rest):
# def pack(pytree, axis=-1):
#   def unpack(array):


from absl.testing import parameterized
from epi_forecast_stat_mech.statistical_models import linear_model
from epi_forecast_stat_mech.statistical_models import tree_util
from absl.testing import absltest
import jax.numpy as jnp
"""Tests for epi_forecast_stat_mech.statistical_models.linear_model."""





class LinearModelTest(parameterized.TestCase):

  def assertBroadcastableTo(self, x, shape):
    try:
      jnp.broadcast_to(x, shape)
    except ValueError:
      self.fail(f'The value `x` with shape {x.shape} '
                f'was not broadcastable to shape {shape}.')

  def get_model(self):
    return linear_model.LinearModel()

  @parameterized.parameters(
      dict(covariate_shapes=(5, 5, 5),
           mech_param_shapes=(2, 3, 4),
           num_epidemics=20),
      dict(covariate_shapes=(1, 2, 3, 4, 5),
           mech_param_shapes=(100,),
           num_epidemics=1000),
      )
  def testShapes(
      self, covariate_shapes, mech_param_shapes, num_epidemics):
    covariate_dim = sum(covariate_shapes)
    prediction_dim = sum(mech_param_shapes)

# Path: epi_forecast_stat_mech/statistical_models/linear_model.py
# class LinearParameters(collections.namedtuple(
#     'LinearModelParameters',
#     ('alpha',  # An array of linear weights.
#      'intercept',  # The intercept of a linear model.
#      ))):
# class LinearModel:
#   def init(cls, covariate_dim, prediction_dim, intercept=None):
#   def log_prior(self, parameters):
#   def log_likelihood(self, parameters, covariates, mechanistic_parameters):
#   def predict(self, parameters, covariates):
# 
# Path: epi_forecast_stat_mech/statistical_models/tree_util.py
# def tree_flatten(pytree):
# def tree_unflatten(treedef, leaves):
# def tree_map(f, structure):
# def tree_multimap(f, structure, *rest):
# def pack(pytree, axis=-1):
#   def unpack(array):


# Path: epi_forecast_stat_mech/statistical_models/linear_model_test.py
from absl.testing import parameterized
from epi_forecast_stat_mech.statistical_models import linear_model
from epi_forecast_stat_mech.statistical_models import tree_util
from absl.testing import absltest
import jax.numpy as jnp

"""Tests for epi_forecast_stat_mech.statistical_models.linear_model."""





class LinearModelTest(parameterized.TestCase):

  def assertBroadcastableTo(self, x, shape):
    try:
      jnp.broadcast_to(x, shape)
    except ValueError:
      self.fail(f'The value `x` with shape {x.shape} '
                f'was not broadcastable to shape {shape}.')

  def get_model(self):
    return linear_model.LinearModel()

  @parameterized.parameters(
      dict(covariate_shapes=(5, 5, 5),
           mech_param_shapes=(2, 3, 4),
           num_epidemics=20),
      dict(covariate_shapes=(1, 2, 3, 4, 5),
           mech_param_shapes=(100,),
           num_epidemics=1000),
      )
  def testShapes(
      self, covariate_shapes, mech_param_shapes, num_epidemics):
    covariate_dim = sum(covariate_shapes)
    prediction_dim = sum(mech_param_shapes)

    params = linear_model.LinearParameters.init(covariate_dim,