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

schema/serde/marketing_data.py

@@ -0,0 +1,1321 @@
+# 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.
+
+"""Serialization and deserialization of `InputData` for Meridian models."""
+
+from collections.abc import Mapping
+import dataclasses
+import datetime as dt
+import functools
+import itertools
+from typing import Sequence
+
+from meridian import constants as c
+from meridian.data import input_data as meridian_input_data
+from mmm.v1.common import date_interval_pb2
+from mmm.v1.marketing import marketing_data_pb2 as marketing_pb
+from schema.serde import constants as sc
+from schema.serde import serde
+from schema.utils import time_record
+import numpy as np
+import xarray as xr
+
+from google.type import date_pb2
+
+# Mapping from DataArray names to coordinate names
+_COORD_NAME_MAP = {
+    c.MEDIA: c.MEDIA_CHANNEL,
+    c.REACH: c.RF_CHANNEL,
+    c.FREQUENCY: c.RF_CHANNEL,
+    c.ORGANIC_MEDIA: c.ORGANIC_MEDIA_CHANNEL,
+    c.ORGANIC_REACH: c.ORGANIC_RF_CHANNEL,
+    c.ORGANIC_FREQUENCY: c.ORGANIC_RF_CHANNEL,
+    c.NON_MEDIA_TREATMENTS: c.NON_MEDIA_CHANNEL,
+}
+
+
+@dataclasses.dataclass(frozen=True)
+class _DeserializedTimeDimension:
+  """Wrapper class for `TimeDimension` proto to provide utility methods during deserialization."""
+
+  _time_dimension: marketing_pb.MarketingDataMetadata.TimeDimension
+
+  def __post_init__(self):
+    if not self._time_dimension.dates:
+      raise ValueError("TimeDimension proto must have at least one date.")
+
+  @functools.cached_property
+  def date_coordinates(self) -> list[dt.date]:
+    """Returns a list of date coordinates in this time dimension."""
+    return [dt.date(d.year, d.month, d.day) for d in self._time_dimension.dates]
+
+  @functools.cached_property
+  def time_dimension_interval(self) -> date_interval_pb2.DateInterval:
+    """Returns the `[start, end)` interval that spans this time dimension.
+
+    This date interval spans all of the date coordinates in this time dimension.
+    """
+    date_intervals = time_record.convert_times_to_date_intervals(
+        self.date_coordinates
+    )
+    return _get_date_interval_from_date_intervals(list(date_intervals.values()))
+
+
+@dataclasses.dataclass(frozen=True)
+class _DeserializedMetadata:
+  """A container for parsed metadata from the `MarketingData` proto.
+
+  Attributes:
+    _metadata: The `MarketingDataMetadata` proto.
+  """
+
+  _metadata: marketing_pb.MarketingDataMetadata
+
+  def __post_init__(self):
+    # Evaluate the properties to trigger validation
+    _ = self.time_dimension
+    _ = self.media_time_dimension
+
+  def _get_time_dimension(self, name: str) -> _DeserializedTimeDimension:
+    """Helper method to get a specific TimeDimension proto by name."""
+    for time_dimension in self._metadata.time_dimensions:
+      if time_dimension.name == name:
+        return _DeserializedTimeDimension(time_dimension)
+    raise ValueError(f"No TimeDimension found with name '{name}' in metadata.")
+
+  @functools.cached_property
+  def time_dimension(self) -> _DeserializedTimeDimension:
+    """Returns the TimeDimension with name 'time'."""
+    return self._get_time_dimension(c.TIME)
+
+  @functools.cached_property
+  def media_time_dimension(self) -> _DeserializedTimeDimension:
+    """Returns the TimeDimension with name 'media_time'."""
+    return self._get_time_dimension(c.MEDIA_TIME)
+
+  @functools.cached_property
+  def channel_dimensions(self) -> Mapping[str, list[str]]:
+    """Returns a mapping of channel dimension names to their corresponding channel coordinate names."""
+    return {
+        cd.name: list(cd.channels) for cd in self._metadata.channel_dimensions
+    }
+
+  @functools.cached_property
+  def channel_types(self) -> Mapping[str, str | None]:
+    """Returns a mapping of individual channel names to their types."""
+    channel_coord_map = {}
+    for name, channels in self.channel_dimensions.items():
+      for channel in channels:
+        channel_coord_map[channel] = _COORD_NAME_MAP.get(
+            name,
+        )
+    return channel_coord_map
+
+
+def _extract_data_array(
+    serialized_data_points: Sequence[marketing_pb.MarketingDataPoint],
+    data_extractor_fn,
+    data_name,
+) -> xr.DataArray | None:
+  """Helper function to extract data into an `xr.DataArray`.
+
+  Args:
+    serialized_data_points: A Sequence of MarketingDataPoint protos.
+    data_extractor_fn: A function that takes a data point and returns either a
+      tuple of `(geo_id, time_str, value)`, or `None` if the data point should
+      be skipped.
+    data_name: The desired name for the `xr.DataArray`.
+
+  Returns:
+    An `xr.DataArray` containing the extracted data, or `None` if no data is
+    found.
+  """
+  data_dict = {}  # (geo_id, time_str) -> value
+  geo_ids = []
+  times = []
+
+  for data_point in serialized_data_points:
+    extraction_result = data_extractor_fn(data_point)
+    if extraction_result is None:
+      continue
+
+    geo_id, time_str, value = extraction_result
+
+    # TODO: Enforce dimension uniqueness in Meridian.
+    if geo_id not in geo_ids:
+      geo_ids.append(geo_id)
+    if time_str not in times:
+      times.append(time_str)
+
+    data_dict[(geo_id, time_str)] = value
+
+  if not data_dict:
+    return None
+
+  data_values = np.array([
+      [data_dict.get((geo_id, time), np.nan) for time in times]
+      for geo_id in geo_ids
+  ])
+
+  return xr.DataArray(
+      data=data_values,
+      coords={
+          c.GEO: geo_ids,
+          c.TIME: times,
+      },
+      dims=(c.GEO, c.TIME),
+      name=data_name,
+  )
+
+
+def _extract_3d_data_array(
+    serialized_data_points: Sequence[marketing_pb.MarketingDataPoint],
+    data_extractor_fn,
+    data_name,
+    third_dim_name,
+    time_dim_name=c.TIME,
+) -> xr.DataArray | None:
+  """Helper function to extract data with 3 dimensions into an `xr.DataArray`.
+
+  The first dimension is always `GEO`, and the second is the time dimension
+  (default: `TIME`).
+
+  Args:
+      serialized_data_points: A sequence of MarketingDataPoint protos.
+      data_extractor_fn: A function that takes a data point and returns either a
+        tuple of `(geo_id, time_str, third_dim_key, value)`, or `None` if the
+        data point should be skipped.
+      data_name: The desired name for the `xr.DataArray`.
+      third_dim_name: The name of the third dimension.
+      time_dim_name: The name of the time dimension. Default is `TIME`.
+
+  Returns:
+      An `xr.DataArray` containing the extracted data, or `None` if no data is
+      found.
+  """
+  data_dict = {}  # (geo_id, time_str, third_dim_key) -> value
+  geo_ids = []
+  times = []
+  third_dim_keys = []
+
+  for data_point in serialized_data_points:
+    for extraction_result in data_extractor_fn(data_point):
+      geo_id, time_str, third_dim_key, value = extraction_result
+
+      if geo_id not in geo_ids:
+        geo_ids.append(geo_id)
+      if time_str not in times:
+        times.append(time_str)
+      if third_dim_key not in third_dim_keys:
+        third_dim_keys.append(third_dim_key)
+
+      # TODO: Enforce dimension uniqueness in Meridian.
+      data_dict[(geo_id, time_str, third_dim_key)] = value
+
+  if not data_dict:
+    return None
+
+  data_values = np.array([
+      [
+          [
+              data_dict.get((geo_id, time, third_dim_key), np.nan)
+              for third_dim_key in third_dim_keys
+          ]
+          for time in times
+      ]
+      for geo_id in geo_ids
+  ])
+
+  return xr.DataArray(
+      data=data_values,
+      coords={
+          c.GEO: geo_ids,
+          time_dim_name: times,
+          third_dim_name: third_dim_keys,
+      },
+      dims=(c.GEO, time_dim_name, third_dim_name),
+      name=data_name,
+  )
+
+
+def _get_date_interval_from_date_intervals(
+    date_intervals: Sequence[date_interval_pb2.DateInterval],
+) -> date_interval_pb2.DateInterval:
+  """Gets the date interval based on the earliest start date and latest end date.
+
+  Args:
+      date_intervals: A list of DateInterval protos.
+
+  Returns:
+      A DateInterval representing the earliest start date and latest end date.
+  """
+  get_start_date = lambda interval: dt.date(
+      interval.start_date.year,
+      interval.start_date.month,
+      interval.start_date.day,
+  )
+  get_end_date = lambda interval: dt.date(
+      interval.end_date.year, interval.end_date.month, interval.end_date.day
+  )
+
+  min_start_date_interval = min(date_intervals, key=get_start_date)
+  max_end_date_interval = max(date_intervals, key=get_end_date)
+
+  return date_interval_pb2.DateInterval(
+      start_date=date_pb2.Date(
+          year=min_start_date_interval.start_date.year,
+          month=min_start_date_interval.start_date.month,
+          day=min_start_date_interval.start_date.day,
+      ),
+      end_date=date_pb2.Date(
+          year=max_end_date_interval.end_date.year,
+          month=max_end_date_interval.end_date.month,
+          day=max_end_date_interval.end_date.day,
+      ),
+  )
+
+
+class _InputDataSerializer:
+  """Serializes an `InputData` container in Meridian model."""
+
+  def __init__(self, input_data: meridian_input_data.InputData):
+    self._input_data = input_data
+
+  @property
+  def _n_geos(self) -> int:
+    return len(self._input_data.geo)
+
+  @property
+  def _n_times(self) -> int:
+    return len(self._input_data.time)
+
+  def __call__(self) -> marketing_pb.MarketingData:
+    """Serializes the input data into a MarketingData proto."""
+    marketing_proto = marketing_pb.MarketingData()
+    # Use media_time since it covers larger range.
+    times_to_date_intervals = time_record.convert_times_to_date_intervals(
+        self._input_data.media_time.data
+    )
+    geos_and_times = itertools.product(
+        self._input_data.geo.data, self._input_data.media_time.data
+    )
+
+    for geo, time in geos_and_times:
+      data_point = self._serialize_data_point(
+          geo,
+          time,
+          times_to_date_intervals,
+      )
+      marketing_proto.marketing_data_points.append(data_point)
+
+    if self._input_data.media_spend is not None:
+      if (
+          not self._input_data.media_spend_has_geo_dimension
+          and not self._input_data.media_spend_has_time_dimension
+      ):
+        marketing_proto.marketing_data_points.append(
+            self._serialize_aggregated_media_spend_data_point(
+                self._input_data.media_spend,
+                times_to_date_intervals,
+            )
+        )
+      elif (
+          self._input_data.media_spend_has_geo_dimension
+          != self._input_data.media_spend_has_time_dimension
+      ):
+        raise AssertionError(
+            "Invalid input data: media_spend must either be fully granular"
+            " (both geo and time dimensions) or fully aggregated (neither geo"
+            " nor time dimensions)."
+        )
+
+    if self._input_data.rf_spend is not None:
+      if (
+          not self._input_data.rf_spend_has_geo_dimension
+          and not self._input_data.rf_spend_has_time_dimension
+      ):
+        marketing_proto.marketing_data_points.append(
+            self._serialize_aggregated_rf_spend_data_point(
+                self._input_data.rf_spend, times_to_date_intervals
+            )
+        )
+      elif (
+          self._input_data.rf_spend_has_geo_dimension
+          != self._input_data.rf_spend_has_time_dimension
+      ):
+        raise AssertionError(
+            "Invalid input data: rf_spend must either be fully granular (both"
+            " geo and time dimensions) or fully aggregated (neither geo nor"
+            " time dimensions)."
+        )
+
+    marketing_proto.metadata.CopyFrom(self._serialize_metadata())
+
+    return marketing_proto
+
+  def _serialize_media_variables(
+      self,
+      geo: str,
+      time: str,
+      channel_dim_name: str,
+      impressions_data_array: xr.DataArray,
+      spend_data_array: xr.DataArray | None = None,
+  ) -> list[marketing_pb.MediaVariable]:
+    """Serializes media variables for a given geo and time.
+
+    Args:
+      geo: The geo ID.
+      time: The time string.
+      channel_dim_name: The name of the channel dimension.
+      impressions_data_array: The DataArray containing impressions data.
+        Expected dimensions: `(n_geos, n_media_times, n_channels)`.
+      spend_data_array: The optional DataArray containing spend data. Expected
+        dimensions are `(n_geos, n_times, n_media_channels)`.
+
+    Returns:
+      A list of MediaVariable protos.
+    """
+    media_variables = []
+    for media_data in impressions_data_array.sel(geo=geo, media_time=time):
+      channel = media_data[channel_dim_name].item()
+      media_variable = marketing_pb.MediaVariable(
+          channel_name=channel,
+          scalar_metric=marketing_pb.ScalarMetric(
+              name=c.IMPRESSIONS, value=media_data.item()
+          ),
+      )
+      if spend_data_array is not None and time in spend_data_array.time:
+        media_variable.media_spend = spend_data_array.sel(
+            geo=geo, time=time, **{channel_dim_name: channel}
+        ).item()
+      media_variables.append(media_variable)
+    return media_variables
+
+  def _serialize_reach_frequency_variables(
+      self,
+      geo: str,
+      time: str,
+      channel_dim_name: str,
+      reach_data_array: xr.DataArray,
+      frequency_data_array: xr.DataArray,
+      spend_data_array: xr.DataArray | None = None,
+  ) -> list[marketing_pb.ReachFrequencyVariable]:
+    """Serializes reach and frequency variables for a given geo and time.
+
+    Iterates through the R&F channels separately, creating a MediaVariable
+    for each. It's safe to assume that Meridian media channel names are
+    unique across `media_data` and `reach_data`. This assumption is
+    checked when an `InputData` is created in model training.
+
+    Dimensions of `reach_data_array` and `frequency_data_array` are expected
+    to be `(n_geos, n_media_times, n_rf_channels)`.
+
+    Args:
+      geo: The geo ID.
+      time: The time string.
+      channel_dim_name: The name of the channel dimension (e.g., 'rf_channel').
+      reach_data_array: The DataArray containing reach data.
+      frequency_data_array: The DataArray containing frequency data.
+      spend_data_array: The optional DataArray containing spend data.
+
+    Returns:
+      A list of MediaVariable protos.
+    """
+    rf_variables = []
+    for reach_data in reach_data_array.sel(geo=geo, media_time=time):
+      reach_value = reach_data.item()
+      channel = reach_data[channel_dim_name].item()
+      frequency_value = frequency_data_array.sel(
+          geo=geo,
+          media_time=time,
+          **{channel_dim_name: channel},
+      ).item()
+      rf_variable = marketing_pb.ReachFrequencyVariable(
+          channel_name=channel,
+          reach=int(reach_value),
+          average_frequency=frequency_value,
+      )
+      if spend_data_array is not None and time in spend_data_array.time:
+        rf_variable.spend = spend_data_array.sel(
+            geo=geo, time=time, **{channel_dim_name: channel}
+        ).item()
+      rf_variables.append(rf_variable)
+    return rf_variables
+
+  def _serialize_non_media_treatment_variables(
+      self, geo: str, time: str
+  ) -> list[marketing_pb.NonMediaTreatmentVariable]:
+    """Serializes non-media treatment variables for a given geo and time.
+
+    Args:
+      geo: The geo ID.
+      time: The time string.
+
+    Returns:
+      A list of NonMediaTreatmentVariable protos.
+    """
+    non_media_treatment_variables = []
+    if (
+        self._input_data.non_media_treatments is not None
+        and geo in self._input_data.non_media_treatments.geo
+        and time in self._input_data.non_media_treatments.time
+    ):
+      for non_media_treatment_data in self._input_data.non_media_treatments.sel(
+          geo=geo, time=time
+      ):
+        non_media_treatment_variables.append(
+            marketing_pb.NonMediaTreatmentVariable(
+                name=non_media_treatment_data[c.NON_MEDIA_CHANNEL].item(),
+                value=non_media_treatment_data.item(),
+            )
+        )
+    return non_media_treatment_variables
+
+  def _serialize_data_point(
+      self,
+      geo: str,
+      time: str,
+      times_to_date_intervals: Mapping[str, date_interval_pb2.DateInterval],
+  ) -> marketing_pb.MarketingDataPoint:
+    """Serializes a MarketingDataPoint proto for a given geo and time."""
+    data_point = marketing_pb.MarketingDataPoint(
+        geo_info=marketing_pb.GeoInfo(
+            geo_id=geo,
+            population=round(self._input_data.population.sel(geo=geo).item()),
+        ),
+        date_interval=times_to_date_intervals.get(time),
+    )
+
+    if self._input_data.controls is not None:
+      if time in self._input_data.controls.time:
+        for control_data in self._input_data.controls.sel(geo=geo, time=time):
+          data_point.control_variables.add(
+              name=control_data.control_variable.item(),
+              value=control_data.item(),
+          )
+
+    if self._input_data.media is not None:
+      if (
+          self._input_data.media_spend_has_geo_dimension
+          and self._input_data.media_spend_has_time_dimension
+      ):
+        spend_data_array = self._input_data.media_spend
+      else:
+        # Aggregated spend data is serialized in a separate data point.
+        spend_data_array = None
+      media_variables = self._serialize_media_variables(
+          geo,
+          time,
+          c.MEDIA_CHANNEL,
+          self._input_data.media,
+          spend_data_array,
+      )
+      data_point.media_variables.extend(media_variables)
+
+    if (
+        self._input_data.reach is not None
+        and self._input_data.frequency is not None
+    ):
+      if (
+          self._input_data.rf_spend_has_geo_dimension
+          and self._input_data.rf_spend_has_time_dimension
+      ):
+        rf_spend_data_array = self._input_data.rf_spend
+      else:
+        # Aggregated spend data is serialized in a separate data point.
+        rf_spend_data_array = None
+      rf_variables = self._serialize_reach_frequency_variables(
+          geo,
+          time,
+          c.RF_CHANNEL,
+          self._input_data.reach,
+          self._input_data.frequency,
+          rf_spend_data_array,
+      )
+      data_point.reach_frequency_variables.extend(rf_variables)
+
+    if self._input_data.organic_media is not None:
+      organic_media_variables = self._serialize_media_variables(
+          geo, time, c.ORGANIC_MEDIA_CHANNEL, self._input_data.organic_media
+      )
+      data_point.media_variables.extend(organic_media_variables)
+
+    if (
+        self._input_data.organic_reach is not None
+        and self._input_data.organic_frequency is not None
+    ):
+      organic_rf_variables = self._serialize_reach_frequency_variables(
+          geo,
+          time,
+          c.ORGANIC_RF_CHANNEL,
+          self._input_data.organic_reach,
+          self._input_data.organic_frequency,
+      )
+      data_point.reach_frequency_variables.extend(organic_rf_variables)
+
+    non_media_treatment_variables = (
+        self._serialize_non_media_treatment_variables(geo, time)
+    )
+    data_point.non_media_treatment_variables.extend(
+        non_media_treatment_variables
+    )
+
+    if time in self._input_data.kpi.time:
+      kpi_proto = self._make_kpi_proto(geo, time)
+      data_point.kpi.CopyFrom(kpi_proto)
+
+    return data_point
+
+  def _serialize_aggregated_media_spend_data_point(
+      self,
+      spend_data_array: xr.DataArray,
+      times_to_date_intervals: Mapping[str, date_interval_pb2.DateInterval],
+  ) -> marketing_pb.MarketingDataPoint:
+    """Serializes and appends a data point for aggregated spend."""
+    spend_data_point = marketing_pb.MarketingDataPoint()
+    date_interval = _get_date_interval_from_date_intervals(
+        list(times_to_date_intervals.values())
+    )
+    spend_data_point.date_interval.CopyFrom(date_interval)
+
+    for channel_name in spend_data_array.coords[c.MEDIA_CHANNEL].values:
+      spend_value = spend_data_array.sel(
+          **{c.MEDIA_CHANNEL: channel_name}
+      ).item()
+      spend_data_point.media_variables.add(
+          channel_name=channel_name, media_spend=spend_value
+      )
+
+    return spend_data_point
+
+  def _serialize_aggregated_rf_spend_data_point(
+      self,
+      spend_data_array: xr.DataArray,
+      times_to_date_intervals: Mapping[str, date_interval_pb2.DateInterval],
+  ) -> marketing_pb.MarketingDataPoint:
+    """Serializes and appends a data point for aggregated spend."""
+    spend_data_point = marketing_pb.MarketingDataPoint()
+    date_interval = _get_date_interval_from_date_intervals(
+        list(times_to_date_intervals.values())
+    )
+    spend_data_point.date_interval.CopyFrom(date_interval)
+
+    for channel_name in spend_data_array.coords[c.RF_CHANNEL].values:
+      spend_value = spend_data_array.sel(**{c.RF_CHANNEL: channel_name}).item()
+      spend_data_point.reach_frequency_variables.add(
+          channel_name=channel_name, spend=spend_value
+      )
+
+    return spend_data_point
+
+  def _serialize_time_dimensions(
+      self, name: str, time_data: xr.DataArray
+  ) -> marketing_pb.MarketingDataMetadata.TimeDimension:
+    """Creates a TimeDimension message."""
+    time_dimensions = marketing_pb.MarketingDataMetadata.TimeDimension(
+        name=name
+    )
+    for date in time_data.values:
+      date_obj = dt.datetime.strptime(date, c.DATE_FORMAT).date()
+      time_dimensions.dates.add(
+          year=date_obj.year, month=date_obj.month, day=date_obj.day
+      )
+    return time_dimensions
+
+  def _serialize_channel_dimensions(
+      self, channel_data: xr.DataArray | None
+  ) -> marketing_pb.MarketingDataMetadata.ChannelDimension | None:
+    """Creates a ChannelDimension message if the corresponding attribute exists."""
+    if channel_data is None:
+      return None
+
+    coord_name = _COORD_NAME_MAP.get(channel_data.name)
+    if coord_name:
+      return marketing_pb.MarketingDataMetadata.ChannelDimension(
+          name=channel_data.name,
+          channels=channel_data.coords[coord_name].values.tolist(),
+      )
+    else:
+      # Make sure that all channel dimensions are handled.
+      raise ValueError(f"Unknown channel data name: {channel_data.name}. ")
+
+  def _serialize_metadata(self) -> marketing_pb.MarketingDataMetadata:
+    """Serializes metadata from InputData to MarketingDataMetadata."""
+    metadata = marketing_pb.MarketingDataMetadata()
+
+    metadata.time_dimensions.append(
+        self._serialize_time_dimensions(c.TIME, self._input_data.time)
+    )
+    metadata.time_dimensions.append(
+        self._serialize_time_dimensions(
+            c.MEDIA_TIME, self._input_data.media_time
+        )
+    )
+
+    channel_data_arrays = [
+        self._input_data.media,
+        self._input_data.reach,
+        self._input_data.frequency,
+        self._input_data.organic_media,
+        self._input_data.organic_reach,
+        self._input_data.organic_frequency,
+    ]
+
+    for channel_data_array in channel_data_arrays:
+      channel_names_message = self._serialize_channel_dimensions(
+          channel_data_array
+      )
+      if channel_names_message:
+        metadata.channel_dimensions.append(channel_names_message)
+
+    if self._input_data.controls is not None:
+      metadata.control_names.extend(
+          self._input_data.controls.control_variable.values
+      )
+
+    if self._input_data.non_media_treatments is not None:
+      metadata.non_media_treatment_names.extend(
+          self._input_data.non_media_treatments.non_media_channel.values
+      )
+
+    metadata.kpi_type = self._input_data.kpi_type
+
+    return metadata
+
+  def _make_kpi_proto(self, geo: str, time: str) -> marketing_pb.Kpi:
+    """Constructs a Kpi proto from the TrainedModel."""
+    kpi_proto = marketing_pb.Kpi(name=self._input_data.kpi_type)
+    # `kpi` and `revenue_per_kpi` dimensions: `(n_geos, n_times)`.
+    if self._input_data.kpi_type == c.REVENUE:
+      kpi_proto.revenue.CopyFrom(
+          marketing_pb.Kpi.Revenue(
+              value=self._input_data.kpi.sel(geo=geo, time=time).item()
+          )
+      )
+    else:
+      kpi_proto.non_revenue.CopyFrom(
+          marketing_pb.Kpi.NonRevenue(
+              value=self._input_data.kpi.sel(geo=geo, time=time).item()
+          )
+      )
+      if self._input_data.revenue_per_kpi is not None:
+        kpi_proto.non_revenue.revenue_per_kpi = (
+            self._input_data.revenue_per_kpi.sel(geo=geo, time=time).item()
+        )
+    return kpi_proto
+
+
+class _InputDataDeserializer:
+  """Deserializes a `MarketingData` proto into a Meridian `InputData`."""
+
+  def __init__(self, serialized: marketing_pb.MarketingData):
+    self._serialized = serialized
+
+  def __post_init__(self):
+    if not self._serialized.HasField(sc.METADATA):
+      raise ValueError(
+          f"MarketingData proto is missing the '{sc.METADATA}' field."
+      )
+
+  @functools.cached_property
+  def _metadata(self) -> _DeserializedMetadata:
+    """Parses metadata and extracts time dimensions, channel dimensions, and channel type map."""
+    return _DeserializedMetadata(self._serialized.metadata)
+
+  def _extract_population(self) -> xr.DataArray:
+    """Extracts population data from the serialized proto."""
+    geo_populations = {}
+
+    for data_point in self._serialized.marketing_data_points:
+      geo_id = data_point.geo_info.geo_id
+      if not geo_id:
+        continue
+
+      geo_populations[geo_id] = data_point.geo_info.population
+
+    return xr.DataArray(
+        coords={c.GEO: list(geo_populations.keys())},
+        data=np.array(list(geo_populations.values())),
+        name=c.POPULATION,
+    )
+
+  def _extract_kpi_type(self) -> str:
+    """Extracts the kpi_type from the serialized proto."""
+    kpi_type = None
+    for data_point in self._serialized.marketing_data_points:
+      if data_point.HasField(c.KPI):
+        current_kpi_type = data_point.kpi.WhichOneof(c.TYPE)
+
+        if kpi_type is None:
+          kpi_type = current_kpi_type
+        elif kpi_type != current_kpi_type:
+          raise ValueError(
+              "Inconsistent kpi_type found in the data. "
+              f"Expected {kpi_type}, found {current_kpi_type}"
+          )
+
+    if kpi_type is None:
+      raise ValueError("kpi_type not found in the data.")
+    return kpi_type
+
+  def _extract_geo_and_time(self, data_point) -> tuple[str | None, str]:
+    """Extracts geo_id and time_str from a data_point."""
+    geo_id = data_point.geo_info.geo_id
+    start_date = data_point.date_interval.start_date
+    time_str = dt.datetime(
+        start_date.year, start_date.month, start_date.day
+    ).strftime(c.DATE_FORMAT)
+    return geo_id, time_str
+
+  def _extract_kpi(self, kpi_type: str) -> xr.DataArray:
+    """Extracts KPI data from the serialized proto."""
+
+    def _kpi_extractor(data_point):
+      if not data_point.HasField(c.KPI):
+        return None
+
+      geo_id, time_str = self._extract_geo_and_time(data_point)
+
+      if data_point.kpi.WhichOneof(c.TYPE) != kpi_type:
+        raise ValueError(
+            "Inconsistent kpi_type found in the data. "
+            f"Expected {kpi_type}, found"
+            f" {data_point.kpi.WhichOneof(c.TYPE)}"
+        )
+
+      kpi_value = (
+          data_point.kpi.revenue.value
+          if kpi_type == c.REVENUE
+          else data_point.kpi.non_revenue.value
+      )
+      return geo_id, time_str, kpi_value
+
+    kpi = _extract_data_array(
+        serialized_data_points=self._serialized.marketing_data_points,
+        data_extractor_fn=_kpi_extractor,
+        data_name=c.KPI,
+    )
+
+    if kpi is None:
+      raise ValueError(f"{c.KPI} is not found in the data.")
+
+    return kpi
+
+  def _extract_revenue_per_kpi(self, kpi_type: str) -> xr.DataArray | None:
+    """Extracts revenue per KPI data from the serialized proto."""
+
+    if kpi_type == c.REVENUE:
+      raise ValueError(
+          f"{c.REVENUE_PER_KPI} is not applicable when kpi_type is {c.REVENUE}."
+      )
+
+    def _revenue_per_kpi_extractor(data_point):
+      if not data_point.HasField(c.KPI):
+        return None
+
+      if not data_point.kpi.non_revenue.HasField(c.REVENUE_PER_KPI):
+        return None
+
+      geo_id, time_str = self._extract_geo_and_time(data_point)
+
+      if data_point.kpi.WhichOneof(c.TYPE) != kpi_type:
+        raise ValueError(
+            "Inconsistent kpi_type found in the data. "
+            f"Expected {kpi_type}, found"
+            f" {data_point.kpi.WhichOneof(c.TYPE)}"
+        )
+
+      return geo_id, time_str, data_point.kpi.non_revenue.revenue_per_kpi
+
+    return _extract_data_array(
+        serialized_data_points=self._serialized.marketing_data_points,
+        data_extractor_fn=_revenue_per_kpi_extractor,
+        data_name=c.REVENUE_PER_KPI,
+    )
+
+  def _extract_controls(self) -> xr.DataArray | None:
+    """Extracts control variables data from the serialized proto, if any."""
+
+    def _controls_extractor(data_point):
+      if not data_point.control_variables:
+        return None
+
+      geo_id, time_str = self._extract_geo_and_time(data_point)
+
+      for control_variable in data_point.control_variables:
+        control_name = control_variable.name
+        control_value = control_variable.value
+        yield geo_id, time_str, control_name, control_value
+
+    return _extract_3d_data_array(
+        serialized_data_points=self._serialized.marketing_data_points,
+        data_extractor_fn=_controls_extractor,
+        data_name=c.CONTROLS,
+        third_dim_name=c.CONTROL_VARIABLE,
+    )
+
+  def _extract_media(self) -> xr.DataArray | None:
+    """Extracts media variables data from the serialized proto."""
+
+    def _media_extractor(data_point):
+      geo_id, time_str = self._extract_geo_and_time(data_point)
+
+      if not geo_id:
+        return None
+
+      for media_variable in data_point.media_variables:
+        channel_name = media_variable.channel_name
+        if self._metadata.channel_types.get(channel_name) != c.MEDIA_CHANNEL:
+          continue
+
+        media_value = media_variable.scalar_metric.value
+        yield geo_id, time_str, channel_name, media_value
+
+    return _extract_3d_data_array(
+        serialized_data_points=self._serialized.marketing_data_points,
+        data_extractor_fn=_media_extractor,
+        data_name=c.MEDIA,
+        third_dim_name=c.MEDIA_CHANNEL,
+        time_dim_name=c.MEDIA_TIME,
+    )
+
+  def _extract_reach(self) -> xr.DataArray | None:
+    """Extracts reach data from the serialized proto."""
+
+    def _reach_extractor(data_point):
+      geo_id, time_str = self._extract_geo_and_time(data_point)
+
+      if not geo_id:
+        return None
+
+      for rf_variable in data_point.reach_frequency_variables:
+        channel_name = rf_variable.channel_name
+        if self._metadata.channel_types.get(channel_name) != c.RF_CHANNEL:
+          continue
+
+        reach_value = rf_variable.reach
+        yield geo_id, time_str, channel_name, reach_value
+
+    return _extract_3d_data_array(
+        serialized_data_points=self._serialized.marketing_data_points,
+        data_extractor_fn=_reach_extractor,
+        data_name=c.REACH,
+        third_dim_name=c.RF_CHANNEL,
+        time_dim_name=c.MEDIA_TIME,
+    )
+
+  def _extract_frequency(self) -> xr.DataArray | None:
+    """Extracts frequency data from the serialized proto."""
+
+    def _frequency_extractor(data_point):
+      geo_id, time_str = self._extract_geo_and_time(data_point)
+
+      if not geo_id:
+        return None
+
+      for rf_variable in data_point.reach_frequency_variables:
+        channel_name = rf_variable.channel_name
+        if self._metadata.channel_types.get(channel_name) != c.RF_CHANNEL:
+          continue
+
+        frequency_value = rf_variable.average_frequency
+        yield geo_id, time_str, channel_name, frequency_value
+
+    return _extract_3d_data_array(
+        serialized_data_points=self._serialized.marketing_data_points,
+        data_extractor_fn=_frequency_extractor,
+        data_name=c.FREQUENCY,
+        third_dim_name=c.RF_CHANNEL,
+        time_dim_name=c.MEDIA_TIME,
+    )
+
+  def _extract_organic_media(self) -> xr.DataArray | None:
+    """Extracts organic media variables data from the serialized proto."""
+
+    def _organic_media_extractor(data_point):
+      geo_id, time_str = self._extract_geo_and_time(data_point)
+
+      if not geo_id:
+        return None
+
+      for media_variable in data_point.media_variables:
+        channel_name = media_variable.channel_name
+        if (
+            self._metadata.channel_types.get(channel_name)
+            != c.ORGANIC_MEDIA_CHANNEL
+        ):
+          continue
+
+        media_value = media_variable.scalar_metric.value
+        yield geo_id, time_str, channel_name, media_value
+
+    return _extract_3d_data_array(
+        serialized_data_points=self._serialized.marketing_data_points,
+        data_extractor_fn=_organic_media_extractor,
+        data_name=c.ORGANIC_MEDIA,
+        third_dim_name=c.ORGANIC_MEDIA_CHANNEL,
+        time_dim_name=c.MEDIA_TIME,
+    )
+
+  def _extract_organic_reach(self) -> xr.DataArray | None:
+    """Extracts organic reach data from the serialized proto."""
+
+    def _organic_reach_extractor(data_point):
+      geo_id, time_str = self._extract_geo_and_time(data_point)
+
+      if not geo_id:
+        return None
+
+      for rf_variable in data_point.reach_frequency_variables:
+        channel_name = rf_variable.channel_name
+        if (
+            self._metadata.channel_types.get(channel_name)
+            != c.ORGANIC_RF_CHANNEL
+        ):
+          continue
+
+        reach_value = rf_variable.reach
+        yield geo_id, time_str, channel_name, reach_value
+
+    return _extract_3d_data_array(
+        serialized_data_points=self._serialized.marketing_data_points,
+        data_extractor_fn=_organic_reach_extractor,
+        data_name=c.ORGANIC_REACH,
+        third_dim_name=c.ORGANIC_RF_CHANNEL,
+        time_dim_name=c.MEDIA_TIME,
+    )
+
+  def _extract_organic_frequency(self) -> xr.DataArray | None:
+    """Extracts organic frequency data from the serialized proto."""
+
+    def _organic_frequency_extractor(data_point):
+      geo_id, time_str = self._extract_geo_and_time(data_point)
+
+      if not geo_id:
+        return None
+
+      for rf_variable in data_point.reach_frequency_variables:
+        channel_name = rf_variable.channel_name
+        if (
+            self._metadata.channel_types.get(channel_name)
+            != c.ORGANIC_RF_CHANNEL
+        ):
+          continue
+
+        frequency_value = rf_variable.average_frequency
+        yield geo_id, time_str, channel_name, frequency_value
+
+    return _extract_3d_data_array(
+        serialized_data_points=self._serialized.marketing_data_points,
+        data_extractor_fn=_organic_frequency_extractor,
+        data_name=c.ORGANIC_FREQUENCY,
+        third_dim_name=c.ORGANIC_RF_CHANNEL,
+        time_dim_name=c.MEDIA_TIME,
+    )
+
+  def _extract_granular_media_spend(
+      self,
+      data_points_with_spend: list[marketing_pb.MarketingDataPoint],
+  ) -> xr.DataArray | None:
+    """Extracts granular spend data.
+
+    Args:
+      data_points_with_spend: List of MarketingDataPoint protos with spend data.
+
+    Returns:
+      An xr.DataArray with granular spend data or None if no data found.
+    """
+
+    def _granular_spend_extractor(data_point):
+      geo_id, time_str = self._extract_geo_and_time(data_point)
+      for media_variable in data_point.media_variables:
+        if (
+            media_variable.HasField(c.MEDIA_SPEND)
+            and self._metadata.channel_types.get(media_variable.channel_name)
+            == c.MEDIA_CHANNEL
+        ):
+          yield geo_id, time_str, media_variable.channel_name, media_variable.media_spend
+
+    return _extract_3d_data_array(
+        serialized_data_points=data_points_with_spend,
+        data_extractor_fn=_granular_spend_extractor,
+        data_name=c.MEDIA_SPEND,
+        third_dim_name=c.MEDIA_CHANNEL,
+        time_dim_name=c.TIME,
+    )
+
+  def _extract_granular_rf_spend(
+      self,
+      data_points_with_spend: list[marketing_pb.MarketingDataPoint],
+  ) -> xr.DataArray | None:
+    """Extracts granular spend data.
+
+    Args:
+      data_points_with_spend: List of MarketingDataPoint protos with spend data.
+
+    Returns:
+      An xr.DataArray with granular spend data or None if no data found.
+    """
+
+    def _granular_spend_extractor(data_point):
+      geo_id, time_str = self._extract_geo_and_time(data_point)
+      for rf_variable in data_point.reach_frequency_variables:
+        if (
+            rf_variable.HasField(c.SPEND)
+            and self._metadata.channel_types.get(rf_variable.channel_name)
+            == c.RF_CHANNEL
+        ):
+          yield geo_id, time_str, rf_variable.channel_name, rf_variable.spend
+
+    return _extract_3d_data_array(
+        serialized_data_points=data_points_with_spend,
+        data_extractor_fn=_granular_spend_extractor,
+        data_name=c.RF_SPEND,
+        third_dim_name=c.RF_CHANNEL,
+        time_dim_name=c.TIME,
+    )
+
+  def _extract_aggregated_media_spend(
+      self,
+      data_points_with_spend: list[marketing_pb.MarketingDataPoint],
+  ) -> xr.DataArray | None:
+    """Extracts aggregated spend data.
+
+    Args:
+      data_points_with_spend: List of MarketingDataPoint protos with spend data.
+
+    Returns:
+      An xr.DataArray with aggregated spend data or None if no data found.
+    """
+    channel_names = self._metadata.channel_dimensions.get(c.MEDIA, [])
+    channel_spend_map = {}
+
+    for spend_data_point in data_points_with_spend:
+      for media_variable in spend_data_point.media_variables:
+        if (
+            media_variable.channel_name in channel_names
+            and media_variable.HasField(c.MEDIA_SPEND)
+        ):
+          channel_spend_map[media_variable.channel_name] = (
+              media_variable.media_spend
+          )
+
+    if not channel_spend_map:
+      return None
+
+    return xr.DataArray(
+        data=list(channel_spend_map.values()),
+        coords={c.MEDIA_CHANNEL: list(channel_spend_map.keys())},
+        dims=[c.MEDIA_CHANNEL],
+        name=c.MEDIA_SPEND,
+    )
+
+  def _extract_aggregated_rf_spend(
+      self,
+      data_points_with_spend: list[marketing_pb.MarketingDataPoint],
+  ) -> xr.DataArray | None:
+    """Extracts aggregated spend data.
+
+    Args:
+      data_points_with_spend: List of MarketingDataPoint protos with spend data.
+
+    Returns:
+      An xr.DataArray with aggregated spend data or None if no data found.
+    """
+    channel_names = self._metadata.channel_dimensions.get(c.REACH, [])
+    channel_spend_map = {}
+
+    for spend_data_point in data_points_with_spend:
+      for rf_variable in spend_data_point.reach_frequency_variables:
+        if rf_variable.channel_name in channel_names and rf_variable.HasField(
+            c.SPEND
+        ):
+          channel_spend_map[rf_variable.channel_name] = rf_variable.spend
+
+    if not channel_spend_map:
+      return None
+
+    return xr.DataArray(
+        data=list(channel_spend_map.values()),
+        coords={c.RF_CHANNEL: list(channel_spend_map.keys())},
+        dims=[c.RF_CHANNEL],
+        name=c.RF_SPEND,
+    )
+
+  def _is_aggregated_spend_data_point(
+      self, dp: marketing_pb.MarketingDataPoint
+  ) -> bool:
+    """Checks if a MarketingDataPoint with spend represents aggregated spend data.
+
+    Args:
+      dp: A marketing_pb.MarketingDataPoint representing a spend data point.
+
+    Returns:
+      True if the data point represents aggregated spend, False otherwise.
+    """
+    if not dp.HasField(sc.GEO_INFO) and self._metadata.media_time_dimension:
+      media_time_interval = (
+          self._metadata.media_time_dimension.time_dimension_interval
+      )
+      return (
+          media_time_interval.start_date == dp.date_interval.start_date
+          and media_time_interval.end_date == dp.date_interval.end_date
+      )
+    return False
+
+  def _extract_media_spend(self) -> xr.DataArray | None:
+    """Extracts media spend data from the serialized proto.
+
+    Returns:
+      An xr.DataArray with spend data or None if no data found.
+    """
+    # Filter data points relevant to spend based on channel type map
+    media_channels = {
+        channel
+        for channel, metadata_channel_type in self._metadata.channel_types.items()
+        if metadata_channel_type == c.MEDIA_CHANNEL
+    }
+    spend_data_points = [
+        dp
+        for dp in self._serialized.marketing_data_points
+        if any(
+            mv.HasField(c.MEDIA_SPEND) and mv.channel_name in media_channels
+            for mv in dp.media_variables
+        )
+    ]
+
+    if not spend_data_points:
+      return None
+
+    aggregated_spend_data_points = [
+        dp
+        for dp in spend_data_points
+        if self._is_aggregated_spend_data_point(dp)
+    ]
+
+    if aggregated_spend_data_points:
+      return self._extract_aggregated_media_spend(aggregated_spend_data_points)
+
+    return self._extract_granular_media_spend(spend_data_points)
+
+  def _extract_rf_spend(self) -> xr.DataArray | None:
+    """Extracts reach and frequency spend data from the serialized proto.
+
+    Returns:
+      An xr.DataArray with spend data or None if no data found.
+    """
+    # Filter data points relevant to spend based on channel type map
+    rf_channels = {
+        channel
+        for channel, metadata_channel_type in self._metadata.channel_types.items()
+        if metadata_channel_type == c.RF_CHANNEL
+    }
+    spend_data_points = [
+        dp
+        for dp in self._serialized.marketing_data_points
+        if any(
+            mv.HasField(c.SPEND) and mv.channel_name in rf_channels
+            for mv in dp.reach_frequency_variables
+        )
+    ]
+
+    if not spend_data_points:
+      return None
+
+    aggregated_spend_data_points = [
+        dp
+        for dp in spend_data_points
+        if self._is_aggregated_spend_data_point(dp)
+    ]
+
+    if aggregated_spend_data_points:
+      return self._extract_aggregated_rf_spend(aggregated_spend_data_points)
+
+    return self._extract_granular_rf_spend(spend_data_points)
+
+  def _extract_non_media_treatments(self) -> xr.DataArray | None:
+    """Extracts non-media treatment variables data from the serialized proto."""
+
+    def _non_media_treatments_extractor(data_point):
+      if not data_point.non_media_treatment_variables:
+        return None
+
+      geo_id, time_str = self._extract_geo_and_time(data_point)
+
+      for (
+          non_media_treatment_variable
+      ) in data_point.non_media_treatment_variables:
+        treatment_name = non_media_treatment_variable.name
+        treatment_value = non_media_treatment_variable.value
+        yield geo_id, time_str, treatment_name, treatment_value
+
+    non_media_treatments_data_array = _extract_3d_data_array(
+        serialized_data_points=self._serialized.marketing_data_points,
+        data_extractor_fn=_non_media_treatments_extractor,
+        data_name=c.NON_MEDIA_TREATMENTS,
+        third_dim_name=c.NON_MEDIA_CHANNEL,
+    )
+
+    return non_media_treatments_data_array
+
+  def __call__(self) -> meridian_input_data.InputData:
+    """Converts the `MarketingData` proto to a Meridian `InputData`."""
+    kpi_type = self._extract_kpi_type()
+    return meridian_input_data.InputData(
+        kpi=self._extract_kpi(kpi_type),
+        kpi_type=kpi_type,
+        controls=self._extract_controls(),
+        population=self._extract_population(),
+        revenue_per_kpi=(
+            self._extract_revenue_per_kpi(kpi_type)
+            if kpi_type == c.NON_REVENUE
+            else None
+        ),
+        media=self._extract_media(),
+        media_spend=self._extract_media_spend(),
+        reach=self._extract_reach(),
+        frequency=self._extract_frequency(),
+        rf_spend=self._extract_rf_spend(),
+        organic_media=self._extract_organic_media(),
+        organic_reach=self._extract_organic_reach(),
+        organic_frequency=self._extract_organic_frequency(),
+        non_media_treatments=self._extract_non_media_treatments(),
+    )
+
+
+class MarketingDataSerde(
+    serde.Serde[marketing_pb.MarketingData, meridian_input_data.InputData]
+):
+  """Serializes and deserializes an `InputData` container in Meridian."""
+
+  def serialize(
+      self, obj: meridian_input_data.InputData
+  ) -> marketing_pb.MarketingData:
+    """Serializes the given Meridian input data into a `MarketingData` proto."""
+    return _InputDataSerializer(obj)()
+
+  def deserialize(
+      self, serialized: marketing_pb.MarketingData, serialized_version: str = ""
+  ) -> meridian_input_data.InputData:
+    """Deserializes the given `MarketingData` proto.
+
+    Args:
+      serialized: The serialized `MarketingData` proto.
+      serialized_version: The version of the serialized model. This is used to
+        handle changes in deserialization logic across different versions.
+
+    Returns:
+      A Meridian input data container.
+    """
+    return _InputDataDeserializer(serialized)()