62 lines
1.7 KiB
Python
62 lines
1.7 KiB
Python
from plotly.subplots import make_subplots
|
|
import plotly.graph_objects as go
|
|
from math import ceil
|
|
from typing import Dict
|
|
import numpy as np
|
|
|
|
|
|
def plot_histograms_in_3d(
|
|
histograms: Dict[str, np.ndarray], histogram_per_row: int = 3
|
|
):
|
|
cols = min(histogram_per_row, len(histograms))
|
|
rows = ceil(len(histograms) / histogram_per_row)
|
|
fig = make_subplots(
|
|
rows=rows,
|
|
cols=cols,
|
|
specs=[[{"type": "scatter3d"} for _ in range(cols)] for _ in range(rows)],
|
|
)
|
|
|
|
for i, (title, histogram) in enumerate(histograms.items()):
|
|
fig.add_trace(
|
|
_get_3d_scatter_plot_from_histogram(title, histogram),
|
|
row=(i // (histogram_per_row + 1)) + 1,
|
|
col=(i % histogram_per_row) + 1,
|
|
)
|
|
|
|
scenes = {
|
|
f"scene{i}": dict(camera=dict(eye=dict(x=0.1, y=0, z=2)))
|
|
for i in range(1, len(histograms) + 1)
|
|
}
|
|
fig.update_layout(**scenes)
|
|
fig.show()
|
|
|
|
|
|
def _get_3d_scatter_plot_from_histogram(title, histogram):
|
|
x, y, z, marker_size = [], [], [], []
|
|
bins = len(histogram)
|
|
|
|
for i, row in enumerate(histogram):
|
|
for j, col in enumerate(row):
|
|
for k, value in enumerate(col):
|
|
if value > 0:
|
|
x.append(i)
|
|
y.append(j)
|
|
z.append(k)
|
|
marker_size.append(value)
|
|
|
|
return go.Scatter3d(
|
|
x=x,
|
|
y=y,
|
|
z=z,
|
|
mode="markers",
|
|
marker=dict(
|
|
size=[min(20, ms * 10000) for ms in marker_size],
|
|
color=[
|
|
f"rgb({xi*256/bins},{yi*256/bins},{zi*256/bins})"
|
|
for xi, yi, zi in zip(x, y, z)
|
|
],
|
|
opacity=1,
|
|
line=dict(width=0),
|
|
),
|
|
name=title,
|
|
)
|