diff --git a/flexmeasures/data/models/generic_assets.py b/flexmeasures/data/models/generic_assets.py index 3ba9330d4b..52ddfffacd 100644 --- a/flexmeasures/data/models/generic_assets.py +++ b/flexmeasures/data/models/generic_assets.py @@ -464,39 +464,78 @@ def chart( for sensor in sensors: sensor.sensor_type = sensor.get_attribute("sensor_type", sensor.name) - # Set up chart specification - if dataset_name is None: - dataset_name = "asset_" + str(self.id) - if event_starts_after: - kwargs["event_starts_after"] = event_starts_after - if event_ends_before: - kwargs["event_ends_before"] = event_ends_before - chart_specs = chart_type_to_chart_specs( + return self._gather_specs_and_data_for_chart( + sensors, chart_type, - sensors_to_show=self.sensors_to_show, - dataset_name=dataset_name, + event_starts_after, + event_ends_before, + beliefs_after, + beliefs_before, + source, + include_data, + dataset_name, + resolution, **kwargs, ) - if include_data: - # Get data - data = self.search_beliefs( - sensors=sensors, - as_json=True, - event_starts_after=event_starts_after, - event_ends_before=event_ends_before, - beliefs_after=beliefs_after, - beliefs_before=beliefs_before, - source=source, - resolution=resolution, - ) + def charts( + self, + chart_type: str = "chart_for_multiple_sensors", + event_starts_after: datetime | None = None, + event_ends_before: datetime | None = None, + beliefs_after: datetime | None = None, + beliefs_before: datetime | None = None, + source: DataSource + | list[DataSource] + | int + | list[int] + | str + | list[str] + | None = None, + include_data: bool = False, + dataset_name: str | None = None, + resolution: str | timedelta | None = None, + **kwargs, + ) -> list[dict]: + """Create multiple vega-lite charts, one chart per entry in `sensors_to_show`. - # Combine chart specs and data - chart_specs["datasets"] = { - dataset_name: json.loads(data), - } + :param chart_type: currently only "bar_chart" # todo: where can we properly list the available chart types? + :param event_starts_after: only return beliefs about events that start after this datetime (inclusive) + :param event_ends_before: only return beliefs about events that end before this datetime (inclusive) + :param beliefs_after: only return beliefs formed after this datetime (inclusive) + :param beliefs_before: only return beliefs formed before this datetime (inclusive) + :param source: search only beliefs by this source (pass the DataSource, or its name or id) or list of sources + :param include_data: if True, include data in the chart, or if False, exclude data + :param dataset_name: optionally name the dataset used in the chart (the default name is sensor_) + :param resolution: optionally set the resolution of data being displayed + :returns: JSON string defining vega-lite chart specs + """ + sensors = flatten_unique(self.sensors_to_show) + for sensor in sensors: + sensor.sensor_type = sensor.get_attribute("sensor_type", sensor.name) - return chart_specs + # Ensure all elements are lists as search_beliefs requires a list of sensors + nested_sensors = [ + [e] if not isinstance(e, list) else e for e in self.sensors_to_show + ] + + return [ + self._gather_specs_and_data_for_chart( + sensors, + chart_type, + event_starts_after, + event_ends_before, + beliefs_after, + beliefs_before, + source, + include_data, + dataset_name, + resolution, + wrap_sensors_in_list=True, + **kwargs, + ) + for sensors in nested_sensors + ] def search_beliefs( self, @@ -754,6 +793,79 @@ def set_inflexible_sensors(self, inflexible_sensor_ids: list[int]) -> None: ).all() db.session.add(self) + def _gather_specs_and_data_for_chart( + self, + sensors: list["Sensor"], # noqa F821 + chart_type: str, + event_starts_after: datetime | None, + event_ends_before: datetime | None, + beliefs_after: datetime | None, + beliefs_before: datetime | None, + source: DataSource + | list[DataSource] + | int + | list[int] + | str + | list[str] + | None, + include_data: bool, + dataset_name: str | None, + resolution: str | timedelta | None, + wrap_sensors_in_list: bool = False, + **kwargs, + ) -> dict: + """ + Centralizes the logic for generating Vega-Lite chart specifications. + + This utility function is designed to be used by both the `chart` and `charts` methods. + + :param sensors: A list of sensors to be displayed in the chart. + :param chart_type: currently only "bar_chart" # todo: where can we properly list the available chart types? + :param event_starts_after: only return beliefs about events that start after this datetime (inclusive) + :param event_ends_before: only return beliefs about events that end before this datetime (inclusive) + :param beliefs_after: only return beliefs formed after this datetime (inclusive) + :param beliefs_before: only return beliefs formed before this datetime (inclusive) + :param source: search only beliefs by this source (pass the DataSource, or its name or id) or list of sources + :param include_data: if True, include data in the chart, or if False, exclude data + :param dataset_name: optionally name the dataset used in the chart (the default name is sensor_) + :param resolution: optionally set the resolution of data being displayed + :param wrap_sensors_in_list: If True, wraps the sensors in a list, which is necessary when displaying multiple sensors in a single chart. + Defaults to False. + :param kwargs: Additional keyword arguments that can be passed to further customize the chart specifications. + :return: A dictionary containing the Vega-Lite chart specification. + """ + + if dataset_name is None: + dataset_name = "asset_" + str(self.id) + if event_starts_after: + kwargs["event_starts_after"] = event_starts_after + if event_ends_before: + kwargs["event_ends_before"] = event_ends_before + + sensors_to_show = [sensors] if wrap_sensors_in_list else self.sensors_to_show + + chart_specs = chart_type_to_chart_specs( + chart_type, + sensors_to_show=sensors_to_show, + dataset_name=dataset_name, + **kwargs, + ) + + if include_data: + data = self.search_beliefs( + sensors=sensors, + as_json=True, + event_starts_after=event_starts_after, + event_ends_before=event_ends_before, + beliefs_after=beliefs_after, + beliefs_before=beliefs_before, + source=source, + resolution=resolution, + ) + chart_specs["datasets"] = {dataset_name: json.loads(data)} + + return chart_specs + def create_generic_asset(generic_asset_type: str, **kwargs) -> GenericAsset: """Create a GenericAsset and assigns it an id.