Skip to content

Fix: Correct GMM fitting visualization in amrex_data.ipynb #165

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
henry2004y wants to merge 2 commits into from
Closed

Fix: Correct GMM fitting visualization in amrex_data.ipynb #165

Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
e87bfe8
Fix: Correct GMM fitting visualization in amrex_data.ipynb
google-labs-jules[bot] Nov 17, 2025
a5d3d3f
Fix: Correct GMM fitting visualization in amrex_data.ipynb
google-labs-jules[bot] Nov 17, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
94 changes: 36 additions & 58 deletions docs/amrex_data.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -665,8 +665,8 @@
"metadata": {},
"outputs": [],
"source": [
"# Since our synthetic data is isotropic, we set isotropic=True to get temperatures\n",
"extracted_temps = AMReXParticleData.get_gmm_temperatures(gmm_synthetic, isotropic=True)"
"# Since our synthetic data is isotropic, we set isotropic=True\n",
"extracted_params = AMReXParticleData.get_gmm_parameters(gmm_synthetic, isotropic=True)"
]
},
{
Expand All @@ -685,11 +685,11 @@
"print(\"--- Original Synthetic Data Parameters ---\")\n",
"print(f\"Population 1: Center = {mean_1.tolist()}, Temperature = {temp_1}\")\n",
"print(f\"Population 2: Center = {mean_2.tolist()}, Temperature = {temp_2}\")\n",
"print(\"\\n--- GMM Extracted Temperatures ---\")\n",
"for i, params in enumerate(extracted_temps):\n",
"print(\"\\n--- GMM Extracted Parameters ---\")\n",
"for i, params in enumerate(extracted_params):\n",
" center = [round(c, 5) for c in params[\"center\"]]\n",
" temp = round(params[\"temperature\"], 5)\n",
" print(f\"Component {i+1}: Center = {center}, Temperature = {temp}\")"
" v_th_sq = round(params[\"v_th_sq\"], 5)\n",
" print(f\"Component {i+1}: Center = {center}, v_th_sq = {v_th_sq}\")"
]
},
{
Expand Down Expand Up @@ -717,7 +717,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"Finally, we visualize the result. The plot below shows a 2D histogram of the synthetic particle data, with the fitted GMM components overlaid as red dashed ellipses. The ellipses represent the 1, 2, and 3-sigma contours of the Gaussian distributions, clearly showing how the GMM has identified the two distinct populations in the data."
"Finally, we visualize the result. The plot below shows a 2D histogram of the synthetic particle data, with the fitted GMM components overlaid as red dashed contours. The contours represent the 1, 2, and 3-sigma contours of the Gaussian distributions, clearly showing how the GMM has identified the two distinct populations in the data."
]
},
{
Expand All @@ -726,7 +726,7 @@
"metadata": {},
"outputs": [],
"source": [
"from matplotlib.patches import Ellipse\n",
"from scipy.stats import multivariate_normal\n",
"\n",
"fig, ax = plt.subplots(figsize=(8, 6), constrained_layout=True)\n",
"\n",
Expand All @@ -738,29 +738,18 @@
"ax.set_xlabel(\"vx\")\n",
"ax.set_ylabel(\"vy\")\n",
"\n",
"# Overlay the GMM ellipses\n",
"# Create a meshgrid for contour plotting\n",
"x = np.linspace(synthetic_data[:, 0].min(), synthetic_data[:, 0].max(), 100)\n",
"y = np.linspace(synthetic_data[:, 1].min(), synthetic_data[:, 1].max(), 100)\n",
"X, Y = np.meshgrid(x, y)\n",
"pos = np.dstack((X, Y))\n",
"\n",
"# Overlay the GMM contours\n",
"for i in range(gmm_synthetic.n_components):\n",
" cov = gmm_synthetic.covariances_[i]\n",
" mean = gmm_synthetic.means_[i]\n",
"\n",
" # Get eigenvalues and eigenvectors to determine ellipse orientation and size\n",
" vals, vecs = np.linalg.eigh(cov)\n",
" angle = np.degrees(np.arctan2(*vecs[:, 0][::-1]))\n",
"\n",
" # Plot ellipses for 1, 2, and 3 standard deviations\n",
" for n_std in [1.0, 2.0, 3.0]:\n",
" width, height = 2 * n_std * np.sqrt(vals)\n",
" ellipse = Ellipse(\n",
" xy=mean,\n",
" width=width,\n",
" height=height,\n",
" angle=angle,\n",
" edgecolor=\"red\",\n",
" facecolor=\"none\",\n",
" lw=1.5,\n",
" ls=\"--\",\n",
" )\n",
" ax.add_patch(ellipse)\n",
" cov = gmm_synthetic.covariances_[i]\n",
" rv = multivariate_normal(mean, cov)\n",
" ax.contour(X, Y, rv.pdf(pos), levels=3, colors='red', linestyles='--')\n",

@gemini-code-assist gemini-code-assist Bot Nov 17, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

The use of levels=3 in ax.contour does not plot 1, 2, and 3-sigma contours as stated in the markdown description. It will instead draw 4 automatically-chosen contour levels. To accurately plot the sigma contours, you should calculate the PDF values corresponding to 1, 2, and 3 standard deviations.

    rv = multivariate_normal(mean, cov)
    levels = rv.pdf(mean) * np.exp(-0.5 * np.array([1.0, 2.0, 3.0])**2)
    ax.contour(X, Y, rv.pdf(pos), levels=np.sort(levels), colors='red', linestyles='--')

"\n",
"ax.set_aspect(\"equal\", \"box\")\n",
"plt.show()"
Expand Down Expand Up @@ -798,8 +787,8 @@
"gmm_anisotropic = GaussianMixture(n_components=1, random_state=0)\n",
"gmm_anisotropic.fit(anisotropic_data)\n",
"\n",
"# Extract temperatures, specifying isotropic=False\n",
"extracted_temps_aniso = AMReXParticleData.get_gmm_temperatures(\n",
"# Extract parameters, specifying isotropic=False\n",
"extracted_params_aniso = AMReXParticleData.get_gmm_parameters(\n",
" gmm_anisotropic, isotropic=False\n",
")\n",
"\n",
Expand All @@ -808,13 +797,13 @@
"print(\n",
" f\"Population 1: Center = {mean_aniso.tolist()}, Temp Parallel = {temp_parallel}, Temp Perp = {temp_perp}\"\n",
")\n",
"print(\"\\n--- GMM Extracted Anisotropic Temperatures ---\")\n",
"for i, params in enumerate(extracted_temps_aniso):\n",
"print(\"\\n--- GMM Extracted Anisotropic Parameters ---\")\n",
"for i, params in enumerate(extracted_params_aniso):\n",
" center = [round(c, 5) for c in params[\"center\"]]\n",
" temp_par_ext = round(params[\"T_parallel\"], 5)\n",
" temp_per_ext = round(params[\"T_perpendicular\"], 5)\n",
" v_th_sq_par_ext = round(params[\"v_th_sq_parallel\"], 5)\n",
" v_th_sq_per_ext = round(params[\"v_th_sq_perpendicular\"], 5)\n",
" print(\n",
" f\"Component {i+1}: Center = {center}, Temp Parallel = {temp_par_ext}, Temp Perp = {temp_per_ext}\"\n",
" f\"Component {i+1}: Center = {center}, v_th_sq_parallel = {v_th_sq_par_ext}, v_th_sq_perp = {v_th_sq_per_ext}\"\n",
" )"
]
},
Expand All @@ -826,37 +815,26 @@
"source": [
"fig, ax = plt.subplots(figsize=(8, 6), constrained_layout=True)\n",
"\n",
"# Plot the 2D histogram of the synthetic anisotropic data\\n\n",
"# Plot the 2D histogram of the synthetic anisotropic data\n",
"ax.hist2d(\n",
" anisotropic_data[:, 0], anisotropic_data[:, 1], bins=50, cmap=\"turbo\", density=True\n",
")\n",
"ax.set_title(\"Anisotropic Synthetic Data with GMM Fit Overlay\")\n",
"ax.set_xlabel(r\"$v_{\\parallel}$\")\n",
"ax.set_ylabel(r\"$v_{\\perp}$\")\n",
"\n",
"# Overlay the GMM ellipses\n",
"# Create a meshgrid for contour plotting\n",
"x = np.linspace(anisotropic_data[:, 0].min(), anisotropic_data[:, 0].max(), 100)\n",
"y = np.linspace(anisotropic_data[:, 1].min(), anisotropic_data[:, 1].max(), 100)\n",
"X, Y = np.meshgrid(x, y)\n",
"pos = np.dstack((X, Y))\n",
"\n",
"# Overlay the GMM contours\n",
"for i in range(gmm_anisotropic.n_components):\n",
" cov = gmm_anisotropic.covariances_[i]\n",
" mean = gmm_anisotropic.means_[i]\n",
"\n",
" # Get eigenvalues and eigenvectors to determine ellipse orientation and size\n",
" vals, vecs = np.linalg.eigh(cov)\n",
" angle = np.degrees(np.arctan2(*vecs[:, 0][::-1]))\n",
"\n",
" # Plot ellipses for 1, 2, and 3 standard deviations\n",
" for n_std in [1.0, 2.0, 3.0]:\n",
" width, height = 2 * n_std * np.sqrt(vals)\n",
" ellipse = Ellipse(\n",
" xy=mean,\n",
" width=width,\n",
" height=height,\n",
" angle=angle,\n",
" edgecolor=\"red\",\n",
" facecolor=\"none\",\n",
" lw=1.5,\n",
" ls=\"--\",\n",
" )\n",
" ax.add_patch(ellipse)\n",
" cov = gmm_anisotropic.covariances_[i]\n",
" rv = multivariate_normal(mean, cov)\n",
" ax.contour(X, Y, rv.pdf(pos), levels=3, colors='red', linestyles='--')\n",

@gemini-code-assist gemini-code-assist Bot Nov 17, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Similar to the isotropic case, using levels=3 here will not produce the desired 1, 2, and 3-sigma contours. It will draw 4 automatically-chosen levels. To correctly visualize the sigma contours, the levels must be calculated based on the probability density function.

    rv = multivariate_normal(mean, cov)
    levels = rv.pdf(mean) * np.exp(-0.5 * np.array([1.0, 2.0, 3.0])**2)
    ax.contour(X, Y, rv.pdf(pos), levels=np.sort(levels), colors='red', linestyles='--')

"\n",
"ax.set_aspect(\"equal\", \"box\")\n",
"plt.show()"
Expand Down