Trame Dashboard

src/python/impactx/dashboard/Analyze/analyzeFunctions.py

@@ -0,0 +1,103 @@
+import pandas as pd
+
+from impactx import distribution
+
+distribution_parameters_file_path = "output_distribution_parameters.txt"
+latticeElement_parameters_file_path = "output_latticeElements_parameters.txt"
+
+
+class analyzeFunctions:
+
+    # -----------------------------------------------------------------------------
+    # Functions for Beam Characteristic and Ref particle data table
+    # -----------------------------------------------------------------------------
+
+    @staticmethod
+    def load_data(file_path):
+        """
+        Function to read provided file_path
+        """
+        df = pd.read_csv(file_path, sep=" ")
+        return df
+
+    @staticmethod
+    def convert_to_dict(combined_data):
+        """
+        Function to convert data into dictionary format.
+        Used to have correct dataType in Vuetify data table.
+        """
+        dictionary = combined_data.to_dict(orient="records")
+        columns = combined_data.columns
+        headers = [
+            {"text": column.strip(), "value": column.strip()} for column in columns
+        ]
+        return dictionary, headers
+
+    @staticmethod
+    def combine_files(file1_name, file2_name):
+        """
+        Function to merge two files together.
+        """
+        file1 = analyzeFunctions.load_data(file1_name)
+        file2 = analyzeFunctions.load_data(file2_name)
+        return pd.merge(file1, file2, how="outer")
+
+    @staticmethod
+    def filter_headers(allHeaders, selected_headers):
+        """
+        Function to retrieve only retrieve
+        user selected headers
+        """
+        filtered_headers = []
+        for selectedHeader in allHeaders:
+            if selectedHeader["value"] in selected_headers:
+                filtered_headers.append(selectedHeader)
+        return filtered_headers
+
+    @staticmethod
+    def filter_data(allData, selected_headers):
+        """
+        Function to retrieve only retrieve data for
+        user selected headers
+        """
+        filtered_data = []
+        for row in allData:
+            filtered_row = {}
+            for key, value in row.items():
+                if key in selected_headers:
+                    filtered_row[key] = value
+            filtered_data.append(filtered_row)
+        return filtered_data
+
+    # -----------------------------------------------------------------------------
+    # Helper functions to read lattice elements and distribution parameter list
+    # -----------------------------------------------------------------------------
+
+    def read_latticeElements_file():
+        """
+        Function to help run impactX simulation
+        Returns a list in correct format to read latticeElements.
+        """
+        file_path = latticeElement_parameters_file_path
+        elements_list = []
+        with open(file_path, "r") as file:
+            lines = file.readlines()
+            for line in lines:
+                line = line.strip()
+                if line.startswith("elements."):
+                    element_code = line.replace("elements.", "").rstrip(",")
+                    elements_list.append(eval(f"elements.{element_code}"))
+        return elements_list
+
+    def read_distribution_file():
+        """
+        Function to help run impactX simulation
+        Reads distribution file path line by line.
+        """
+        file_path = distribution_parameters_file_path
+        safe_env = {"distribution": distribution, "distr": None}
+
+        with open(file_path, "r") as file:
+            exec(file.read(), safe_env)
+
+        return safe_env["distr"]

src/python/impactx/dashboard/Analyze/plot_over_s/overS.py

@@ -0,0 +1,32 @@
+import plotly.graph_objects as go
+
+
+def line_plot_1d(selected_headers, filtered_data):
+    x_axis = selected_headers[1] if len(selected_headers) > 1 else None
+    y_axis = selected_headers[2:] if len(selected_headers) > 2 else None
+
+    x = [row[x_axis] for row in filtered_data] if x_axis else []
+
+    figure_data = []
+    if y_axis:
+        for column in y_axis:
+            y = [row[column] for row in filtered_data]
+            trace = go.Scatter(
+                x=x,
+                y=y,
+                mode="lines+markers",
+                name=column,
+                line=dict(width=2),
+                marker=dict(size=8),
+            )
+            figure_data.append(trace)
+
+    return go.Figure(
+        data=figure_data,
+        layout=go.Layout(
+            title="Plot Over S",
+            xaxis=dict(title="s"),
+            yaxis=dict(title=""),
+            margin=dict(l=20, r=20, t=25, b=30),
+        ),
+    )

src/python/impactx/dashboard/Analyze/plot_phase_space/oldHolderSettings.py

@@ -0,0 +1,132 @@
+### Only the single plot shows deltaX_delxaP_X
+
+
+def adjusted_settings_plot(pc, num_bins=50, root_rank=0):
+    import matplotlib.pyplot as plt
+    import numpy as np
+    from scipy.stats import gaussian_kde
+
+    # Data Histogramming
+    df = pc.to_df(local=True)
+
+    if df is None:
+        print("DataFrame is None")
+        return None
+
+    # Update for plot unit system
+    m2mm = 1.0e3
+    rad2mrad = 1.0e3
+
+    # Convert positions to mm and momenta to mrad
+    df.position_x = df.position_x.multiply(m2mm)
+    df.momentum_x = df.momentum_x.multiply(rad2mrad)
+
+    # Calculate the density of points using a 2D histogram
+    xy = np.vstack([df.position_x, df.momentum_x])
+    z = gaussian_kde(xy)(xy)
+
+    # Simple scatter plot of position_x and momentum_x with density coloring
+    fig, ax = plt.subplots(figsize=(8, 4), constrained_layout=True)
+    scatter = ax.scatter(df.position_x, df.momentum_x, c=z, cmap="viridis", alpha=0.5)
+    ax.set_xlabel("Delta x [mm]")
+    ax.set_ylabel("Delta p_x [mrad]")
+
+    # Adding a color bar
+    cbar = plt.colorbar(scatter, ax=ax)
+    cbar.set_label("Particle Density")
+
+    # Adding a legend for the scatter plot (optional if needed for more information)
+    # We can create custom legend if needed, but color bar usually serves the purpose for density.
+
+    """
+    # Beam Characteristics
+    rbc = pc.reduced_beam_characteristics()
+
+    # calculate local histograms
+    xpx, x_edges, px_edges = np.histogram2d(
+        df["position_x"],
+        df["momentum_x"],
+        bins=num_bins,
+        range=[
+            [rbc["x_min"] * m2mm, rbc["x_max"] * m2mm],
+            [rbc["px_min"] * rad2mrad, rbc["px_max"] * rad2mrad],
+        ],
+    )
+
+    # histograms per axis
+    x = np.sum(xpx, axis=1)
+    px = np.sum(xpx, axis=0)
+
+    # Check histogram contents
+    print(f"xpx:\n{xpx}\nx:\n{x}\npx:\n{px}\n")
+
+    # Matplotlib canvas: figure and plottable axes areas
+    fig, ax_xpx = plt.subplots(figsize=(8, 4), constrained_layout=True)
+
+    #   projected axes
+    ax_x, ax_px = ax_xpx.twinx(), ax_xpx.twiny()
+
+    def plot_2d(hist, r, p, r_edges, p_edges, ax_r, ax_p, ax_rp):
+        hist = np.ma.masked_where(hist == 0, hist)
+        im = ax_rp.imshow(
+            hist.T,
+            origin="lower",
+            aspect="auto",
+            extent=[r_edges[0], r_edges[-1], p_edges[0], p_edges[-1]],
+        )
+        cbar = fig.colorbar(im, ax=ax_rp)
+
+        r_mids = (r_edges[:-1] + r_edges[1:]) / 2
+        p_mids = (p_edges[:-1] + p_edges[1:]) / 2
+        ax_r.plot(r_mids, r, c="w", lw=0.8, alpha=0.7)
+        ax_r.plot(r_mids, r, c="k", lw=0.5, alpha=0.7)
+        ax_r.fill_between(r_mids, r, facecolor="k", alpha=0.2)
+        ax_p.plot(p, p_mids, c="w", lw=0.8, alpha=0.7)
+        ax_p.plot(p, p_mids, c="k", lw=0.5, alpha=0.7)
+        ax_p.fill_betweenx(p_mids, p, facecolor="k", alpha=0.2)
+
+        return cbar
+
+    cbar_xpx = plot_2d(xpx, x, px, x_edges, px_edges, ax_x, ax_px, ax_xpx)
+
+    def set_limits(r, p, r_edges, p_edges, ax_r, ax_p, ax_rp):
+        pad = 0.1
+        len_r = r_edges[-1] - r_edges[0]
+        len_p = p_edges[-1] - p_edges[0]
+        ax_rp.set_xlim(r_edges[0] - len_r * pad, r_edges[-1] + len_r * pad)
+        ax_rp.set_ylim(p_edges[0] - len_p * pad, p_edges[-1] + len_p * pad)
+
+        def on_xlims_change(axes):
+            if not axes.xlim_reset_in_progress:
+                pad = 6.0
+                axes.xlim_reset_in_progress = True
+                axes.set_xlim(0, np.max(p) * pad)
+                axes.xlim_reset_in_progress = False
+
+        ax_p.xlim_reset_in_progress = False
+        ax_p.callbacks.connect("xlim_changed", on_xlims_change)
+        on_xlims_change(ax_p)
+
+        def on_ylims_change(axes):
+            if not axes.ylim_reset_in_progress:
+                pad = 6.0
+                axes.ylim_reset_in_progress = True
+                axes.set_ylim(0, np.max(r) * pad)
+                axes.ylim_reset_in_progress = False
+
+        ax_r.ylim_reset_in_progress = False
+        ax_r.callbacks.connect("ylim_changed", on_ylims_change)
+        on_ylims_change(ax_r)
+
+    set_limits(x, px, x_edges, px_edges, ax_x, ax_px, ax_xpx)
+
+    ax_xpx.set_xlabel("Delta x [mm]")
+    ax_xpx.set_ylabel("Delta p_x [mrad]")
+    cbar_xpx.set_label("Q [C/bin]")
+    ax_x.set_yticks([])
+    ax_px.set_xticks([])
+
+    return fig
+    """
+
+    return fig

src/python/impactx/dashboard/Analyze/plot_phase_space/phaseSpace.py

@@ -0,0 +1,99 @@
+#!/usr/bin/env python3
+#
+# Copyright 2022-2023 The ImpactX Community
+#
+# Authors: Axel Huebl
+# License: BSD-3-Clause-LBNL
+#
+# -*- coding: utf-8 -*-
+
+from trame.app import get_server
+
+server = get_server(client_type="vue2")
+state, ctrl = server.state, server.controller
+import importlib
+
+import pytest
+from Analyze.analyzeFunctions import analyzeFunctions
+from Analyze.plot_phase_space.phaseSpaceSettings import adjusted_settings_plot
+
+from impactx import ImpactX
+
+
+@pytest.mark.skipif(
+    importlib.util.find_spec("pandas") is None, reason="pandas is not available"
+)
+def run_simulation(save_png=True):
+    """
+    This tests using ImpactX and Pandas Dataframes
+    """
+    sim = ImpactX()
+
+    sim.particle_shape = state.particle_shape
+    sim.space_charge = False
+    sim.slice_step_diagnostics = False
+    sim.init_grids()
+
+    # init particle beam
+    kin_energy_MeV = state.kin_energy_MeV
+    bunch_charge_C = state.bunch_charge_C
+    npart = state.npart
+
+    #   reference particle
+    pc = sim.particle_container()
+    ref = pc.ref_particle()
+    ref.set_charge_qe(-1.0).set_mass_MeV(0.510998950).set_kin_energy_MeV(kin_energy_MeV)
+
+    #   particle bunch
+    # distr = distribution.Waterbag(
+    #     lambdaX=3.9984884770e-5,
+    #     lambdaY=3.9984884770e-5,
+    #     lambdaT=1.0e-3,
+    #     lambdaPx=2.6623538760e-5,
+    #     lambdaPy=2.6623538760e-5,
+    #     lambdaPt=2.0e-3,
+    #     muxpx=-0.846574929020762,
+    #     muypy=0.846574929020762,
+    #     mutpt=0.0,
+    # )
+    distr = analyzeFunctions.read_distribution_file()
+    sim.add_particles(bunch_charge_C, distr, npart)
+
+    assert pc.total_number_of_particles() == npart
+
+    # init accelerator lattice
+    # fodo = [
+    #     elements.Drift(0.25),
+    #     elements.Quad(1.0, 1.0),
+    #     elements.Drift(0.5),
+    #     elements.Quad(1.0, -1.0),
+    #     elements.Drift(0.25),
+    # ]
+    fodo = analyzeFunctions.read_latticeElements_file()
+
+    sim.lattice.extend(fodo)
+
+    # simulate
+    sim.evolve()
+
+    # check local particles
+    df = pc.to_df(local=True)
+    print(df)
+
+    # ensure the column heads are correctly labeled
+    assert df.columns.tolist() == [
+        "idcpu",
+        "position_x",
+        "position_y",
+        "position_t",
+        "momentum_x",
+        "momentum_y",
+        "momentum_t",
+        "qm",
+        "weighting",
+    ]
+
+    # fig = pc.plot_phasespace()
+    fig = adjusted_settings_plot(pc)
+
+    return fig