feat: created plots for talk to ipcc

climateqa/engine/talk_to_data/ipcc/plots.py

@@ -0,0 +1,193 @@
+from typing import Callable
+from plotly.graph_objects import Figure
+import plotly.graph_objects as go
+import pandas as pd
+import geojson
+
+from climateqa.engine.talk_to_data.ipcc.config import IPCC_INDICATOR_TO_UNIT, IPCC_SCENARIO
+from climateqa.engine.talk_to_data.ipcc.queries import indicator_for_given_year_query, indicator_per_year_at_location_query
+from climateqa.engine.talk_to_data.objects.plot import Plot
+
+def plot_indicator_evolution_at_location_historical_and_projections(
+    params: dict,
+) -> Callable[[pd.DataFrame], Figure]:
+    """
+    Returns a function that generates a line plot showing the evolution of a climate indicator
+    (e.g., temperature, rainfall) over time at a specific location, including both historical data
+    and future projections for different climate scenarios.
+
+    Args:
+        params (dict): Dictionary with:
+            - indicator_column (str): Name of the climate indicator column to plot.
+            - location (str): Location (e.g., country, city) for which to plot the indicator.
+
+    Returns:
+        Callable[[pd.DataFrame], Figure]: Function that takes a DataFrame and returns a Plotly Figure
+        showing the indicator's evolution over time, with scenario lines and historical data.
+    """
+    indicator = params["indicator_column"]
+    location = params["location"]
+    indicator_label = " ".join(word.capitalize() for word in indicator.split("_"))
+    unit = IPCC_INDICATOR_TO_UNIT.get(indicator, "")
+
+    def plot_data(df: pd.DataFrame) -> Figure:
+        df = df.sort_values(by='year')
+        years = df['year'].astype(int).tolist()
+        indicators = df[indicator].astype(float).tolist()
+        scenarios = df['scenario'].astype(str).tolist()
+
+        # Find last historical value for continuity
+        last_historical = [(y, v) for y, v, s in zip(years, indicators, scenarios) if s == 'historical']
+        last_historical_year, last_historical_indicator = last_historical[-1] if last_historical else (None, None)
+
+        fig = go.Figure()
+        for scenario in IPCC_SCENARIO:
+            x = [y for y, s in zip(years, scenarios) if s == scenario]
+            y = [v for v, s in zip(indicators, scenarios) if s == scenario]
+            # Connect historical to scenario
+            if scenario != 'historical' and last_historical_indicator is not None:
+                x = [last_historical_year] + x
+                y = [last_historical_indicator] + y
+            fig.add_trace(go.Scatter(
+                x=x,
+                y=y,
+                mode='lines',
+                name=scenario
+            ))
+
+        fig.update_layout(
+            title=f'Yearly Evolution of {indicator_label} in {location} (Historical + SSP Scenarios)',
+            xaxis_title='Year',
+            yaxis_title=f'{indicator_label} ({unit})',
+            legend_title='Scenario',
+            height=800,
+        )
+        return fig
+
+    return plot_data
+
+indicator_evolution_at_location_historical_and_projections: Plot = {
+    "name": "Indicator Evolution at Location (Historical + Projections)",
+    "description": (
+        "Shows how a climate indicator (e.g., rainfall, temperature) changes over time at a specific location, "
+        "including historical data and future projections. "
+        "Useful for questions about the value or trend of an indicator at a location for any year, "
+        "such as 'What will be the total rainfall in China in 2050?' or 'How does rainfall evolve in China over time?'. "
+        "Parameters: indicator_column (the climate variable), location (e.g., country, city)."
+    ),
+    "params": ["indicator_column", "location"],
+    "plot_function": plot_indicator_evolution_at_location_historical_and_projections,
+    "sql_query": indicator_per_year_at_location_query,
+    "short_name": "Indicator Evolution"
+}
+
+def plot_choropleth_map_of_country_indicator_for_specific_year(
+    params: dict,
+) -> Callable[[pd.DataFrame], Figure]:
+    """
+    Returns a function that generates a choropleth map (heatmap) showing the spatial distribution
+    of a climate indicator (e.g., temperature, rainfall) across all regions of a country for a specific year.
+
+    Args:
+        params (dict): Dictionary with:
+            - indicator_column (str): Name of the climate indicator column to plot.
+            - year (str or int, optional): Year for which to plot the indicator (default: 2050).
+            - country_name (str): Name of the country.
+            - location (str): Location (country or region) for the map.
+
+    Returns:
+        Callable[[pd.DataFrame], Figure]: Function that takes a DataFrame and returns a Plotly Figure
+        showing the indicator's spatial distribution as a choropleth map for the specified year.
+    """
+    indicator = params["indicator_column"]
+    year = params.get('year', 2050)
+    country_name = params['country_name']
+    location = params['location']
+    indicator_label = " ".join(word.capitalize() for word in indicator.split("_"))
+    unit = IPCC_INDICATOR_TO_UNIT.get(indicator, "")
+
+    def plot_data(df: pd.DataFrame) -> Figure:
+        custom_colorscale = [
+            [0.0, "rgb(5, 48, 97)"],
+            [0.10, "rgb(33, 102, 172)"],
+            [0.20, "rgb(67, 147, 195)"],
+            [0.30, "rgb(146, 197, 222)"],
+            [0.40, "rgb(209, 229, 240)"],
+            [0.50, "rgb(247, 247, 247)"],
+            [0.60, "rgb(253, 219, 199)"],
+            [0.75, "rgb(244, 165, 130)"],
+            [0.85, "rgb(214, 96, 77)"],
+            [0.90, "rgb(178, 24, 43)"],
+            [1.0, "rgb(103, 0, 31)"]
+        ]
+
+        indicators = df[indicator].astype(float).tolist()
+        latitudes = df["latitude"].astype(float).tolist()
+        longitudes = df["longitude"].astype(float).tolist()
+
+        features = [
+            geojson.Feature(
+                geometry=geojson.Polygon([[
+                    [lon - 0.5, lat - 0.5],
+                    [lon + 0.5, lat - 0.5],
+                    [lon + 0.5, lat + 0.5],
+                    [lon - 0.5, lat + 0.5],
+                    [lon - 0.5, lat - 0.5]
+                ]]),
+                properties={"value": val},
+                id=str(idx)
+            )
+            for idx, (lat, lon, val) in enumerate(zip(latitudes, longitudes, indicators))
+        ]
+
+        geojson_data = geojson.FeatureCollection(features)
+
+        fig = go.Figure(go.Choroplethmapbox(
+            geojson=geojson_data,
+            locations=[str(i) for i in range(len(indicators))],
+            featureidkey="id",
+            z=indicators,
+            colorscale=custom_colorscale,
+            zmin=min(indicators),
+            zmax=max(indicators),
+            marker_opacity=0.7,
+            marker_line_width=0,
+            colorbar_title=f"{indicator_label} ({unit})",
+            text=[f"{indicator_label}: {value:.2f} {unit}" for value in indicators],  # Add hover text showing the indicator value
+            hoverinfo="text" 
+        ))
+
+        fig.update_layout(
+            mapbox_style="open-street-map",
+            mapbox_zoom=3,
+            height=800,
+            mapbox_center={
+                "lat": latitudes[len(latitudes)//2] if latitudes else 0,
+                "lon": longitudes[len(longitudes)//2] if longitudes else 0
+            },
+            coloraxis_colorbar=dict(title=f"{indicator_label} ({unit})"),
+            title=f"{indicator_label} in {year} in {location} ({country_name})"
+        )
+        return fig
+
+    return plot_data
+
+choropleth_map_of_country_indicator_for_specific_year: Plot = {
+    "name": "Choropleth Map of a Country's Indicator Distribution for a Specific Year",
+    "description": (
+        "Displays a map showing the spatial distribution of a climate indicator (e.g., rainfall, temperature) "
+        "across all regions of a country for a specific year. "
+        "Can answer questions about the value of an indicator in a country or region for a given year, "
+        "such as 'What will be the total rainfall in China in 2050?' or 'How is rainfall distributed across China in 2050?'. "
+        "Parameters: indicator_column (the climate variable), year, location (country name)."
+    ),
+    "params": ["indicator_column", "year", "location"],
+    "plot_function": plot_choropleth_map_of_country_indicator_for_specific_year,
+    "sql_query": indicator_for_given_year_query,
+    "short_name": "Choropleth Map"
+}
+
+IPCC_PLOTS = [
+    indicator_evolution_at_location_historical_and_projections,
+    choropleth_map_of_country_indicator_for_specific_year
+]