google-meridian 1.3.0__py3-none-any.whl → 1.3.1__py3-none-any.whl

meridian/model/eda/eda_engine.py

@@ -175,10 +175,22 @@ def _data_array_like(
   )
 
 
-def _stack_variables(
+def stack_variables(
     ds: xr.Dataset, coord_name: str = _STACK_VAR_COORD_NAME
 ) -> xr.DataArray:
-  """Stacks data variables other than time and geo into a single variable."""
+  """Stacks data variables of a Dataset into a single DataArray.
+
+  This function is designed to work with Datasets that have 'time' or 'geo'
+  dimensions, which are preserved. Other dimensions are stacked into a new
+  dimension.
+
+  Args:
+    ds: The input xarray.Dataset to stack.
+    coord_name: The name of the new coordinate for the stacked dimension.
+
+  Returns:
+    An xarray.DataArray with the specified dimensions stacked.
+  """
   dims = []
   coords = []
   sample_dims = []
@@ -423,10 +435,11 @@ class EDAEngine:
   @functools.cached_property
   def national_media_spend_da(self) -> xr.DataArray | None:
     """Returns the national media spend data array."""
-    if self.media_spend_da is None:
+    media_spend = self.media_spend_da
+    if media_spend is None:
       return None
     if self._is_national_data:
-      national_da = self.media_spend_da.squeeze(constants.GEO, drop=True)
+      national_da = media_spend.squeeze(constants.GEO, drop=True)
       national_da.name = constants.NATIONAL_MEDIA_SPEND
     else:
       national_da = self._aggregate_and_scale_geo_da(
@@ -576,10 +589,11 @@ class EDAEngine:
   @functools.cached_property
   def national_rf_spend_da(self) -> xr.DataArray | None:
     """Returns the national RF spend data array."""
-    if self.rf_spend_da is None:
+    rf_spend = self.rf_spend_da
+    if rf_spend is None:
       return None
     if self._is_national_data:
-      national_da = self.rf_spend_da.squeeze(constants.GEO, drop=True)
+      national_da = rf_spend.squeeze(constants.GEO, drop=True)
       national_da.name = constants.NATIONAL_RF_SPEND
     else:
       national_da = self._aggregate_and_scale_geo_da(
@@ -806,10 +820,42 @@ class EDAEngine:
     ]
     return xr.merge(to_merge, join='inner')
 
+  @functools.cached_property
+  def all_spend_ds(self) -> xr.Dataset:
+    """Returns a Dataset containing all spend data.
+
+    This includes media spend and rf spend.
+    """
+    to_merge = [
+        da
+        for da in [
+            self.media_spend_da,
+            self.rf_spend_da,
+        ]
+        if da is not None
+    ]
+    return xr.merge(to_merge, join='inner')
+
+  @functools.cached_property
+  def national_all_spend_ds(self) -> xr.Dataset:
+    """Returns a Dataset containing all national spend data.
+
+    This includes media spend and rf spend.
+    """
+    to_merge = [
+        da
+        for da in [
+            self.national_media_spend_da,
+            self.national_rf_spend_da,
+        ]
+        if da is not None
+    ]
+    return xr.merge(to_merge, join='inner')
+
   @functools.cached_property
   def _stacked_treatment_control_scaled_da(self) -> xr.DataArray:
     """Returns a stacked DataArray of treatment_control_scaled_ds."""
-    da = _stack_variables(self.treatment_control_scaled_ds)
+    da = stack_variables(self.treatment_control_scaled_ds)
     da.name = constants.TREATMENT_CONTROL_SCALED
     return da
 
@@ -837,7 +883,7 @@ class EDAEngine:
   @functools.cached_property
   def _stacked_national_treatment_control_scaled_da(self) -> xr.DataArray:
     """Returns a stacked DataArray of national_treatment_control_scaled_ds."""
-    da = _stack_variables(self.national_treatment_control_scaled_ds)
+    da = stack_variables(self.national_treatment_control_scaled_ds)
     da.name = constants.NATIONAL_TREATMENT_CONTROL_SCALED
     return da
 

meridian/model/model_test_data.py

@@ -143,6 +143,7 @@ class WithInputDataSamples:
   _short_input_data_with_rf_only: input_data.InputData
   _short_input_data_with_media_and_rf: input_data.InputData
   _national_input_data_media_only: input_data.InputData
+  _national_input_data_rf_only: input_data.InputData
   _national_input_data_media_and_rf: input_data.InputData
   _test_dist_media_and_rf: collections.OrderedDict[str, backend.Tensor]
   _test_dist_media_only: collections.OrderedDict[str, backend.Tensor]
@@ -282,6 +283,16 @@ class WithInputDataSamples:
             seed=0,
         )
     )
+    cls._national_input_data_rf_only = (
+        test_utils.sample_input_data_non_revenue_revenue_per_kpi(
+            n_geos=cls._N_GEOS_NATIONAL,
+            n_times=cls._N_TIMES,
+            n_media_times=cls._N_MEDIA_TIMES,
+            n_controls=cls._N_CONTROLS,
+            n_rf_channels=cls._N_RF_CHANNELS,
+            seed=0,
+        )
+    )
     cls._national_input_data_media_only = (
         test_utils.sample_input_data_non_revenue_revenue_per_kpi(
             n_geos=cls._N_GEOS_NATIONAL,
@@ -581,6 +592,10 @@ class WithInputDataSamples:
   def national_input_data_media_only(self) -> input_data.InputData:
     return self._national_input_data_media_only.copy(deep=True)
 
+  @property
+  def national_input_data_rf_only(self) -> input_data.InputData:
+    return self._national_input_data_rf_only.copy(deep=True)
+
   @property
   def national_input_data_media_and_rf(self) -> input_data.InputData:
     return self._national_input_data_media_and_rf.copy(deep=True)

meridian/version.py

@@ -14,4 +14,4 @@
 
 """Module for Meridian version."""
 
-__version__ = "1.3.0"
+__version__ = "1.3.1"

schema/__init__.py

@@ -0,0 +1,18 @@
+# Copyright 2025 The Meridian Authors.
+#
+# 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 containing MMM schema library."""
+
+from schema import serde
+from schema import utils

schema/serde/__init__.py

@@ -0,0 +1,26 @@
+# Copyright 2025 The Meridian Authors.
+#
+# 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.
+
+"""A serialization and deserialization library for Meridian models.
+
+For entry points API, see `meridian_serde` module docs.
+"""
+
+from schema.serde import constants
+from schema.serde import distribution
+from schema.serde import eda_spec
+from schema.serde import hyperparameters
+from schema.serde import inference_data
+from schema.serde import meridian_serde
+from schema.serde import serde

schema/serde/constants.py

@@ -0,0 +1,48 @@
+# Copyright 2025 The Meridian Authors.
+#
+# 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.
+
+"""Constants shared across the Meridian serde library."""
+
+# Constants for hyperparameters protobuf structure
+BASELINE_GEO_ONEOF = 'baseline_geo_oneof'
+BASELINE_GEO_INT = 'baseline_geo_int'
+BASELINE_GEO_STRING = 'baseline_geo_string'
+CONTROL_POPULATION_SCALING_ID = 'control_population_scaling_id'
+HOLDOUT_ID = 'holdout_id'
+NON_MEDIA_POPULATION_SCALING_ID = 'non_media_population_scaling_id'
+ADSTOCK_DECAY_SPEC = 'adstock_decay_spec'
+GLOBAL_ADSTOCK_DECAY = 'global_adstock_decay'
+ADSTOCK_DECAY_BY_CHANNEL = 'adstock_decay_by_channel'
+DEFAULT_DECAY = 'geometric'
+
+# Constants for marketing data protobuf structure
+GEO_INFO = 'geo_info'
+METADATA = 'metadata'
+REACH_FREQUENCY = 'reach_frequency'
+
+# Constants for distribution protobuf structure
+DISTRIBUTION_TYPE = 'distribution_type'
+BATCH_BROADCAST_DISTRIBUTION = 'batch_broadcast'
+DETERMINISTIC_DISTRIBUTION = 'deterministic'
+HALF_NORMAL_DISTRIBUTION = 'half_normal'
+LOG_NORMAL_DISTRIBUTION = 'log_normal'
+NORMAL_DISTRIBUTION = 'normal'
+TRANSFORMED_DISTRIBUTION = 'transformed'
+TRUNCATED_NORMAL_DISTRIBUTION = 'truncated_normal'
+UNIFORM_DISTRIBUTION = 'uniform'
+BETA_DISTRIBUTION = 'beta'
+BIJECTOR_TYPE = 'bijector_type'
+SHIFT_BIJECTOR = 'shift'
+SCALE_BIJECTOR = 'scale'
+RECIPROCAL_BIJECTOR = 'reciprocal'