google-meridian 1.1.0__py3-none-any.whl → 1.1.2__py3-none-any.whl

meridian/data/nd_array_input_data_builder.py

@@ -0,0 +1,509 @@
+# 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.
+
+"""An implementation of `InputDataBuilder` with n-dimensional array primitives."""
+
+import logging
+import warnings
+from meridian import constants
+from meridian.data import input_data_builder
+import numpy as np
+import xarray as xr
+
+
+__all__ = [
+    'NDArrayInputDataBuilder',
+]
+
+
+class NDArrayInputDataBuilder(input_data_builder.InputDataBuilder):
+  """Builds `InputData` from n-dimensional arrays."""
+
+  # Unlike `DataFrameInputDataBuilder`, each piecemeal data has no coordinate
+  # information; they're purely data values. It's up to the user to provide
+  # coordinates with setter methods from the abstract base class above.
+  # Validation is done on each piece w.r.t. dimensional consistency by
+  # shape alone.
+
+  def with_kpi(self, nd: np.ndarray) -> 'NDArrayInputDataBuilder':
+    """Reads KPI data from a ndarray.
+
+    `nd` must be given with the shape:
+    -  `(n_geos, n_time)`
+    -  `(n_time,)` or `(1, n_time)` for national model.
+
+    If called without a call to .geos() first, the data will be
+    assumed to be national-level.
+
+    Args:
+      nd: The ndarray to read the KPI data from.
+
+    Returns:
+      The `NDArrayInputDataBuilder` with the added KPI data.
+    """
+    ### Validate ###
+    self._validate_coords()
+    self._validate_shape(nd)
+
+    ### Transform ###
+    self.kpi = xr.DataArray(
+        nd,
+        dims=[constants.GEO, constants.TIME],
+        coords={
+            constants.GEO: self.geos,
+            constants.TIME: self.time_coords,
+        },
+        name=constants.KPI,
+    )
+    return self
+
+  def with_controls(
+      self, nd: np.ndarray, control_names: list[str]
+  ) -> 'NDArrayInputDataBuilder':
+    """Reads controls data from a ndarray.
+
+    `nd` must be given with the shape:
+    -  `(n_geos, n_time, n_controls)`
+    -  `(n_time, n_controls)` or `(1, n_time, n_controls)` for national model.
+
+    If called without a call to .geos() first, the data will be
+    assumed to be national-level.
+
+    Args:
+      nd: The ndarray to read the controls data from.
+      control_names: The names of the control variables.
+
+    Returns:
+      The `NDArrayInputDataBuilder` with the added controls data.
+    """
+    ### Validate ###
+    self._validate_coords()
+    self._validate_shape(nd, control_names)
+
+    ### Transform ###
+    self.controls = xr.DataArray(
+        nd,
+        dims=[constants.GEO, constants.TIME, constants.CONTROL_VARIABLE],
+        coords={
+            constants.GEO: self.geos,
+            constants.TIME: self.time_coords,
+            constants.CONTROL_VARIABLE: control_names,
+        },
+        name=constants.CONTROLS,
+    )
+    return self
+
+  def with_population(self, nd: np.ndarray) -> 'NDArrayInputDataBuilder':
+    """Reads population data from a ndarray.
+
+    `nd` must be given with the shape:
+    -  `(n_geos,)`
+
+    If called without a call to .geos() first, the data will be
+    assumed to be national-level.
+
+    Args:
+      nd: The ndarray to read the population data from.
+
+    Returns:
+      The `NDArrayInputDataBuilder` with the added population data.
+    """
+    ### Validate ###
+    self._validate_coords(is_population=True)
+    self._validate_shape(nd, is_population=True)
+    ### Transform ###
+    self.population = xr.DataArray(
+        nd,
+        dims=[constants.GEO],
+        coords={constants.GEO: self.geos},
+        name=constants.POPULATION,
+    )
+
+    return self
+
+  def with_revenue_per_kpi(self, nd: np.ndarray) -> 'NDArrayInputDataBuilder':
+    """Reads Revenue per KPI data from a ndarray.
+
+    `nd` must be given with the shape:
+    -  `(n_geos, n_time)`
+    -  `(n_time,)` or `(1, n_time)` for national model.
+
+    If called without a call to .geos() first, the data will be
+    assumed to be national-level.
+
+    Args:
+      nd: The ndarray to read the Reevenue per KPI data from.
+
+    Returns:
+      The `NDArrayInputDataBuilder` with the added Revenue per KPI data.
+    """
+    ### Validate ###
+    self._validate_coords()
+    self._validate_shape(nd)
+    revenue_per_kpi_nd = self._check_revenue_per_kpi_defaults(nd)
+
+    ### Transform ###
+    self.revenue_per_kpi = xr.DataArray(
+        revenue_per_kpi_nd,
+        dims=[constants.GEO, constants.TIME],
+        coords={
+            constants.GEO: self.geos,
+            constants.TIME: self.time_coords,
+        },
+        name=constants.REVENUE_PER_KPI,
+    )
+    return self
+
+  def with_media(
+      self, m_nd: np.ndarray, ms_nd: np.ndarray, media_channels: list[str]
+  ) -> 'NDArrayInputDataBuilder':
+    """Reads media and media spend data from the ndarrays.
+
+    `m_nd` must be given with the shape:
+    -  `(n_geos, n_media_times, n_media_channels)`
+    -  `(n_media_times, n_media_channels)` or `(1, n_media_times,
+       n_media_channels)` for national model.
+
+    `ms_nd` must be given with the shape:
+    -  `(n_geos, n_times, n_media_channels)`
+    -  `(n_times, n_media_channels)` or `(1, n_times,
+       n_media_channels)` for national model.
+
+    If called without a call to .geos() first, the data will be
+    assumed to be national-level.
+
+    Args:
+      m_nd: The ndarray that contains dimensional media data.
+      ms_nd: The ndarray that contains dimensional media spend data.
+      media_channels: The names of the media channels.
+
+    Returns:
+      The `NDArrayInputDataBuilder` with the added media and media spend data.
+    """
+    ### Validate ###
+    self._validate_coords(is_media_time=True)
+    self._validate_coords(is_media_time=False)
+    self._validate_shape(nd=m_nd, dims=media_channels, is_media_time=True)
+    self._validate_shape(nd=ms_nd, dims=media_channels, is_media_time=False)
+
+    ### Transform ###
+    self.media = xr.DataArray(
+        m_nd,
+        dims=[
+            constants.GEO,
+            constants.MEDIA_TIME,
+            constants.MEDIA_CHANNEL,
+        ],
+        coords={
+            constants.GEO: self.geos,
+            constants.MEDIA_TIME: self.media_time_coords,
+            constants.MEDIA_CHANNEL: media_channels,
+        },
+        name=constants.MEDIA,
+    )
+    self.media_spend = xr.DataArray(
+        ms_nd,
+        dims=[
+            constants.GEO,
+            constants.TIME,
+            constants.MEDIA_CHANNEL,
+        ],
+        coords={
+            constants.GEO: self.geos,
+            constants.TIME: self.time_coords,
+            constants.MEDIA_CHANNEL: media_channels,
+        },
+        name=constants.MEDIA_SPEND,
+    )
+    return self
+
+  def with_reach(
+      self,
+      r_nd: np.ndarray,
+      f_nd: np.ndarray,
+      rfs_nd: np.ndarray,
+      rf_channels: list[str],
+  ) -> 'NDArrayInputDataBuilder':
+    """Reads reach, frequency, and rf_spend data from the ndarrays.
+
+    `r_nd` and `f_nd` must be given with the shape:
+    -  `(n_geos, n_media_times, n_rf_channels)`
+    -  `(n_media_times, n_rf_channels)` or `(1, n_media_times,
+       n_rf_channels)` for national model.
+
+    `rfs_nd` must be given with the shape:
+    -  `(n_geos, n_times, n_rf_channels)`
+    -  `(n_times, n_rf_channels)` or `(1, n_times,
+       n_rf_channels)` for national model.
+
+    If called without a call to .geos() first, the data will be
+    assumed to be national-level.
+
+    Args:
+      r_nd: The ndarray that contains dimensional reach data.
+      f_nd: The ndarray that contains dimensional frequency data.
+      rfs_nd: The ndarray that contains dimensional rf_spend data.
+      rf_channels: The names of the rf channels.
+
+    Returns:
+      The `NDArrayInputDataBuilder` with the added reach, frequency, and
+      rf_spend data.
+    """
+    ### Validate ###
+    self._validate_coords(is_media_time=True)
+    self._validate_coords(is_media_time=False)
+    self._validate_shape(nd=r_nd, dims=rf_channels, is_media_time=True)
+    self._validate_shape(nd=f_nd, dims=rf_channels, is_media_time=True)
+    self._validate_shape(nd=rfs_nd, dims=rf_channels, is_media_time=False)
+
+    ### Transform ###
+    self.reach = xr.DataArray(
+        r_nd,
+        dims=[
+            constants.GEO,
+            constants.MEDIA_TIME,
+            constants.RF_CHANNEL,
+        ],
+        coords={
+            constants.GEO: self.geos,
+            constants.MEDIA_TIME: self.media_time_coords,
+            constants.RF_CHANNEL: rf_channels,
+        },
+        name=constants.REACH,
+    )
+    self.frequency = xr.DataArray(
+        f_nd,
+        dims=[
+            constants.GEO,
+            constants.MEDIA_TIME,
+            constants.RF_CHANNEL,
+        ],
+        coords={
+            constants.GEO: self.geos,
+            constants.MEDIA_TIME: self.media_time_coords,
+            constants.RF_CHANNEL: rf_channels,
+        },
+        name=constants.FREQUENCY,
+    )
+    self.rf_spend = xr.DataArray(
+        rfs_nd,
+        dims=[
+            constants.GEO,
+            constants.TIME,
+            constants.RF_CHANNEL,
+        ],
+        coords={
+            constants.GEO: self.geos,
+            constants.TIME: self.time_coords,
+            constants.RF_CHANNEL: rf_channels,
+        },
+        name=constants.RF_SPEND,
+    )
+    return self
+
+  def with_organic_media(
+      self, nd: np.ndarray, organic_media_channels: list[str]
+  ) -> 'NDArrayInputDataBuilder':
+    """Reads organic media data from a ndarray.
+
+    `nd` must be given with the shape:
+    -  `(n_geos, n_media_times, n_organic_media_channels)`
+    -  `(n_media_times, n_organic_media_channels)` or `(1, n_media_times,
+       n_organic_media_channels)` for national model.
+
+    If called without a call to .geos() first, the data will be
+    assumed to be national-level.
+
+    Args:
+      nd: The ndarray to read the organic media data from.
+      organic_media_channels: The names of the organic media channels.
+
+    Returns:
+      The `NDArrayInputDataBuilder` with the added organic media data.
+    """
+    ### Validate ###
+    self._validate_coords(is_media_time=True)
+    self._validate_shape(nd=nd, dims=organic_media_channels, is_media_time=True)
+
+    ### Transform ###
+    self.organic_media = xr.DataArray(
+        nd,
+        dims=[
+            constants.GEO,
+            constants.MEDIA_TIME,
+            constants.ORGANIC_MEDIA_CHANNEL,
+        ],
+        coords={
+            constants.GEO: self.geos,
+            constants.MEDIA_TIME: self.media_time_coords,
+            constants.ORGANIC_MEDIA_CHANNEL: organic_media_channels,
+        },
+        name=constants.ORGANIC_MEDIA,
+    )
+    return self
+
+  def with_organic_reach(
+      self, or_nd: np.ndarray, of_nd: np.ndarray, organic_rf_channels: list[str]
+  ) -> 'NDArrayInputDataBuilder':
+    """Reads organic reach and organic frequency data from the ndarrays.
+
+    `or_nd` and `of_nd` must be given with the shape:
+    -  `(n_geos, n_media_times, n_organic_rf_channels)`
+    -  `(n_media_times, n_organic_rf_channels)` or `(1, n_media_times,
+       n_organic_rf_channels)` for national model.
+
+    If called without a call to .geos() first, the data will be
+    assumed to be national-level.
+
+    Args:
+      or_nd: The ndarray that contains dimensional reach data.
+      of_nd: The ndarray that contains dimensional frequency data.
+      organic_rf_channels: The names of the organic rf channels.
+
+    Returns:
+      The `NDArrayInputDataBuilder` with the added organic reach and organic
+      frequency data.
+    """
+    ### Validate ###
+    self._validate_coords(is_media_time=True)
+    self._validate_shape(nd=or_nd, dims=organic_rf_channels, is_media_time=True)
+    self._validate_shape(nd=of_nd, dims=organic_rf_channels, is_media_time=True)
+
+    ### Transform ###
+    self.organic_reach = xr.DataArray(
+        or_nd,
+        dims=[
+            constants.GEO,
+            constants.MEDIA_TIME,
+            constants.ORGANIC_RF_CHANNEL,
+        ],
+        coords={
+            constants.GEO: self.geos,
+            constants.MEDIA_TIME: self.media_time_coords,
+            constants.ORGANIC_RF_CHANNEL: organic_rf_channels,
+        },
+        name=constants.ORGANIC_REACH,
+    )
+    self.organic_frequency = xr.DataArray(
+        of_nd,
+        dims=[
+            constants.GEO,
+            constants.MEDIA_TIME,
+            constants.ORGANIC_RF_CHANNEL,
+        ],
+        coords={
+            constants.GEO: self.geos,
+            constants.MEDIA_TIME: self.media_time_coords,
+            constants.ORGANIC_RF_CHANNEL: organic_rf_channels,
+        },
+        name=constants.ORGANIC_REACH,
+    )
+    return self
+
+  def with_non_media_treatments(
+      self, nd: np.ndarray, non_media_channel_names: list[str]
+  ) -> 'NDArrayInputDataBuilder':
+    """Reads non-media treatments data from a ndarray.
+
+    `nd` must be given with the shape:
+    -  `(n_geos, n_time, n_media_channels)`
+    -  `(n_time, n_media_channels)` or `(1, n_time, n_media_channels)` for
+    national model.
+
+    If called without a call to .geos() first, the data will be
+    assumed to be national-level.
+
+    Args:
+      nd: The ndarray to read the non-media treatments data from.
+      non_media_channel_names: The names of the non-media channels.
+
+    Returns:
+      The `NDArrayInputDataBuilder` with the added non-media treatments data.
+    """
+    ### Validate ###
+    self._validate_coords()
+    self._validate_shape(nd, non_media_channel_names)
+
+    ### Transform ###
+    self.non_media_treatments = xr.DataArray(
+        nd,
+        dims=[constants.GEO, constants.TIME, constants.NON_MEDIA_CHANNEL],
+        coords={
+            constants.GEO: self.geos,
+            constants.TIME: self.time_coords,
+            constants.NON_MEDIA_CHANNEL: non_media_channel_names,
+        },
+        name=constants.NON_MEDIA_TREATMENTS,
+    )
+    return self
+
+  def _validate_coords(
+      self, is_population: bool = False, is_media_time: bool = False
+  ):
+    """Validates that the data has the expected coordinates."""
+    if not is_population:
+      if is_media_time and self._media_time_coords is None:
+        raise ValueError(
+            'Media times are required first. Set using .media_time_coords()'
+        )
+      if not is_media_time and self.time_coords is None:
+        raise ValueError(
+            'Time coordinates are required first. Set using .time_coords()'
+        )
+    if self.geos is None:
+      logging.warning(
+          'No geo coordinates set. Assuming NATIONAL model and geos will be set'
+          ' to the default value.'
+      )
+      self.geos = [constants.NATIONAL_MODEL_DEFAULT_GEO_NAME]
+
+  def _validate_shape(
+      self,
+      nd: np.ndarray,
+      dims: list[str] | None = None,
+      is_population: bool = False,
+      is_media_time: bool = False,
+  ):
+    """Validates that the data has the expected shape."""
+    # Since all data has a geo dimension (even for national data),
+    # Expect the first axis to have the shape of the geo dimension.
+    expected_shape = (len(self.geos),)
+    detailed_info = f'Expected: {len(self.geos)} geos'
+    if not is_population:
+      if is_media_time:
+        expected_shape += (len(self.media_time_coords),)
+        detailed_info += f' x {len(self.media_time_coords)} media times'
+      else:
+        expected_shape += (len(self.time_coords),)
+        detailed_info += f' x {len(self.time_coords)} times'
+
+      if dims is not None:
+        if len(dims) != len(set(dims)):
+          raise ValueError('given dimensions must be unique.')
+        expected_shape += (len(dims),)
+        detailed_info += f' x {len(dims)} dims'
+
+    if expected_shape != nd.shape:
+      raise ValueError(f'{detailed_info}. Got: {nd.shape}.')
+
+  def _check_revenue_per_kpi_defaults(self, nd: np.ndarray):
+    """Sets revenue_per_kpi to default if kpi type is revenue and with_revenue_per_kpi is called."""
+    if self._kpi_type == constants.REVENUE:
+      warnings.warn(
+          'with_revenue_per_kpi was called but kpi_type was set to revenue.'
+          ' Assuming revenue per kpi with values [1].'
+      )
+      return np.ones(nd.shape)
+    else:
+      return nd