px_utilities.py

import plotly.graph_objects as go
import numpy as np
import stretchmap_utilities as su
import sham_utilities
from plotly.subplots import make_subplots
import hydrostatic as hy

style_target = "dash"
density_color = "blue"
energy_color = "orange"
pressure_color = "green"
soundspeed_color = "magenta"

i = 0
density_i = i
i += 2
energy_i = i
i += 1
soundspeed_i = i
i += 2
pressure_i = i
i += 2
scatter3d_i = i
dotsize = 8

# import plotly.io as pio

# pio.templates.default = "plotly_dark"


def get_rho_values(data, mpart):
    H = data["hpart"]
    Rho = mpart * (1.2 / H) ** 3
    return Rho


def format_param(key, value):
    if key in ["y0", "n"]:
        valuestr = value
    elif key in ["mtot_target"]:
        valuestr = f"{value:.3f}"
    elif isinstance(value, float):
        if value > 10 or value < 0.1:
            valuestr = f"{value:.1e}"
        else:
            valuestr = f"{value:.1f}"
    else:
        valuestr = value
    return valuestr


def format_inputparams(input_params):
    string = ""
    for key, value in input_params.items():
        if isinstance(value, dict):
            line = ""
            line += f"{key}[name]={format_param("name",value["name"])} <br>"
            for param, param_value in value["values"].items():
                line += f"{key}[{param}]={format_param(param,param_value)} <br>"
        else:
            valuestr = format_param(key, value)
            line = f"{key}: {valuestr} <br>"
        string += line
    return string


def px_3d_and_rho(model, ctx, img_path, rhotarget, input_params, unit):
    """
    I don't normalize here, please be sure that rhotarget is normalized

    :param model: Description
    :param ctx: Description
    :param img_path: Description
    :param rhotarget: Description
    :param input_params: Description
    """
    eos = input_params["eos"]
    data = ctx.collect_data()
    t = model.get_time()
    mpart = model.get_particle_mass()
    Pos = data["xyz"]
    U_int = data["uint"]

    R = np.linalg.norm(Pos, axis=1)
    Np = Pos.shape[0]

    Rho = get_rho_values(data, mpart)

    haspressure = False
    if "pressure" in data:
        haspressure = True
        pressure_data = data["pressure"]
    # haspressure = False

    nrows = 2
    ncols = 2
    fig = make_subplots(
        rows=nrows,
        cols=ncols,
        specs=[
            [{"secondary_y": True}, {"type": "scatter3d", "rowspan": 2}],
            [{"secondary_y": haspressure}, {}],
        ],
        horizontal_spacing=0.2,
        vertical_spacing=0.05,
        column_widths=[0.6, 0.4],
        subplot_titles=(
            "Density and Energy profiles",
            "Soundspeed and Pressure profiles",
            "",
        ),
    )

    rhotab = rhotarget[1]
    rtab = rhotarget[0]
    # ! Density profile
    fig.add_trace(
        go.Scatter(
            x=R,
            y=Rho,
            mode="markers",
            name="rhoSPH",
            marker=dict(color=density_color, size=dotsize),
            opacity=0.6,
        ),
        row=1,
        col=1,
        secondary_y=False,
    )
    fig.add_trace(
        go.Scatter(
            x=rtab,
            y=rhotab,
            mode="lines",
            name="Initial target",
            line=dict(color=density_color, dash=style_target),
        ),
        row=1,
        col=1,
        secondary_y=False,
    )
    fig.add_trace(
        go.Scatter(
            x=R,
            y=U_int,
            mode="markers",
            name="uint",
            marker=dict(color=energy_color, size=dotsize),
            opacity=0.6,
        ),
        row=1,
        col=1,
        secondary_y=True,
    )

    P_cs_func = su.get_p_and_cs_func(eos, unit)
    mask = rhotab != 0
    P_target, cs_target = P_cs_func(rhotab[mask])

    cs_data = data["soundspeed"]

    # ! Soundspeed profile
    fig.add_trace(
        go.Scatter(
            x=R,
            y=cs_data,
            mode="markers",
            marker=dict(color=soundspeed_color, size=dotsize),
            name="soundspeed",
            opacity=0.6,
        ),
        row=2,
        col=1,
        secondary_y=False,
    )

    fig.add_trace(
        go.Scatter(
            x=rtab[mask],
            y=cs_target,
            mode="lines",
            line=dict(color=soundspeed_color, dash=style_target),
            showlegend=False,
        ),
        row=2,
        col=1,
        secondary_y=False,
    )

    # ! Pressure profile
    if haspressure:
        fig.add_trace(
            go.Scatter(
                x=R,
                y=pressure_data,
                mode="markers",
                marker=dict(color=pressure_color, size=dotsize),
                name="pressure",
                opacity=0.6,
            ),
            row=2,
            col=1,
            secondary_y=True,
        )
        fig.add_trace(
            go.Scatter(
                x=rtab[mask],
                y=P_target,
                mode="lines",
                line=dict(color=pressure_color, dash=style_target),
                showlegend=False,
            ),
            row=2,
            col=1,
            secondary_y=True,
        )

    # ! 3D scatter
    fig.add_trace(
        go.Scatter3d(
            x=Pos[:, 0],
            y=Pos[:, 1],
            z=Pos[:, 2],
            mode="markers",
            name="",
            marker=dict(
                color=Rho,
                colorscale="Viridis",
                colorbar=dict(
                    title="Density",
                    tickfont=dict(size=24),
                    xanchor="left",
                    exponentformat="e",
                ),
                cmin=0,
                cmax=np.max(Rho),
                opacity=0.3,
                size=4,
            ),
            showlegend=False,
        ),
        col=2,
        row=1,
    )

    fig.update_xaxes(title_text=r"$r / R_\odot$", row=1, col=1)

    for pos in ([1, 1], [2, 1]):
        row, col = pos
        fig.update_xaxes(
            row=row,
            col=col,
            title_font=dict(size=20),
            tickfont=dict(size=16),
            exponentformat="e",
        )
    fig.update_xaxes(row=1, col=1, title_text=r"$r/R_\odot$")
    fig.update_xaxes(row=2, col=1, title_text=r"$r/R_\odot$")

    fig.update_yaxes(
        row=1,
        col=1,
        title_text=r"$\rho$",
        title_font=dict(color=density_color),
        tickfont=dict(color=density_color),
        secondary_y=False,
    )
    fig.update_yaxes(
        row=1,
        col=1,
        title_text=r"$u$",
        title_font=dict(color=energy_color),
        tickfont=dict(color=energy_color),
        secondary_y=True,
    )
    fig.update_yaxes(
        row=2,
        col=1,
        title_text=r"$c_s$",
        title_font=dict(color=soundspeed_color),
        tickfont=dict(color=soundspeed_color),
    )
    if haspressure:
        fig.update_yaxes(
            row=2,
            col=1,
            title_text=r"$P$",
            title_font=dict(color=pressure_color),
            tickfont=dict(color=pressure_color),
            secondary_y=True,
        )
    fig.update_layout(
        scene=dict(
            xaxis=dict(title=""),
            yaxis=dict(title=""),
            zaxis=dict(title=""),
        )
    )

    fig.update_layout(
        height=860,
        width=1920,
        title_text=f"{img_path} t={t:.2e}",
        annotations=[
            dict(
                text=f"{Np} particles, {Np*mpart:.1e} solar masses",
                showarrow=False,
                xref="paper",
                yref="paper",
                x=0.5,
                y=1.03,
                xanchor="center",
                yanchor="bottom",
                font=dict(size=12, color="gray"),
            ),
            dict(
                text=format_inputparams(input_params),
                align="left",
                showarrow=False,
                xref="paper",
                yref="paper",
                x=0.65,
                y=0.1,
                bordercolor="black",
                borderwidth=1,
            ),
        ],
        font=dict(family="Courier New, monospace"),
        legend=dict(
            orientation="h",
            yanchor="bottom",
            y=1.02,
            xanchor="center",
            x=0.25,
        ),
    )

    # fig.add_annotation(
    #     text="Density profile",
    #     xref="paper",
    #     yref="paper",
    #     x=0,  # left column center (approx)
    #     y=0.98,
    #     showarrow=False,
    #     font=dict(size=16),
    # )
    # fig.add_annotation(
    #     text="Soundspeed and Pressure profiles",
    #     xref="paper",
    #     yref="paper",
    #     x=0,  # right/top area (approx) adjust as required
    #     y=0.5,
    #     showarrow=False,
    #     font=dict(size=16),
    # )

    return fig


def update_px_3d_and_rho(fig, model, ctx, img_path, input_params):
    data = ctx.collect_data()
    t = model.get_time()
    mpart = model.get_particle_mass()

    Pos = data["xyz"]
    U_int = data["uint"]
    R = np.linalg.norm(Pos, axis=1)
    Rho = get_rho_values(data, mpart)

    Np = Pos.shape[0]
    mtot_target = input_params["mtot_target"]

    fig.data[density_i].update(
        x=R,
        y=Rho,
    )

    fig.data[energy_i].update(
        x=R,
        y=U_int,
    )

    cs_data = data["soundspeed"]
    fig.data[soundspeed_i].update(
        x=R,
        y=cs_data,
    )

    fig.update_yaxes(range=[0, None], row=1, col=1, secondary_y=False)
    fig.update_yaxes(range=[0, None], row=1, col=1, secondary_y=True)
    fig.update_yaxes(range=[0, None], row=2, col=1, secondary_y=False)

    if "pressure" in data:
        pressure_data = data["pressure"]
        print("pressure", pressure_data)
        fig.data[pressure_i].update(
            x=R,
            y=pressure_data,
        )
        fig.update_yaxes(
            range=[0, None],
            row=2,
            col=1,
            secondary_y=True,
        )

    # Trace 2 : Scatter3D
    fig.data[scatter3d_i].update(
        x=Pos[:, 0],
        y=Pos[:, 1],
        z=Pos[:, 2],
        marker_color=Rho,
        # marker=dict(
        #     cmin=0,
        #     cmax=np.max(Rho),
        # ),
    )

    fig.update_yaxes(range=[0, 1.1 * np.max(Rho)], row=1, col=1)
    fig.update_layout(
        title_text=f"{img_path} t={t:.2e}",
        annotations=[
            dict(
                text=f"{Np} particles, {Np*mpart:.1e} solar masses",
                showarrow=False,
                xref="paper",
                yref="paper",
                x=0.5,
                y=1.03,
                xanchor="center",
                yanchor="bottom",
                font=dict(size=12, color="gray"),
            ),
            dict(
                text=format_inputparams(input_params),
                align="left",
                showarrow=False,
                xref="paper",
                yref="paper",
                x=0.55,
                y=0.5,
                bordercolor="black",
                borderwidth=1,
            ),
        ],
    )

    return fig


def movie(pattern_png, filemp4, fps):
    import ffmpeg

    ffmpeg.input(pattern_png, pattern_type="glob", framerate=fps).output(
        filemp4,
        vcodec="libx264",
        crf=18,
        preset="medium",
        r=fps,
        pix_fmt="yuv420p",
        movflags="faststart",
    ).overwrite_output().run()


def compute_fps(inputparams):
    """
    so that the duration of the movie is 4 sec
    """
    nb_dumps = inputparams["nb_dumps"]
    # tf = inputparams["tf"]
    return int(nb_dumps / 4)