Skip to content

Commit 3f1b8bd

Browse files
nvaytetnvaytet
authored
Merge pull request #90 from scipp/mcstas-to-nexus
Mcstas to nexus
2 parents bb3e4de + d5b935d commit 3f1b8bd

File tree

1 file changed

+272
-0
lines changed

1 file changed

+272
-0
lines changed

‎tools/mcstas_to_nexus.ipynb‎

Lines changed: 272 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,272 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "markdown",
5+
"id": "0",
6+
"metadata": {},
7+
"source": [
8+
"# McStas to NeXus for ODIN\n",
9+
"\n",
10+
"This notebook converts data from a McStas simulation output (`.h5` filetype) to a NeXus file that uses a file for the Odin instrument (written by CODA) as a template for the geometry information.\n",
11+
"\n",
12+
"It adds events to the `timepix3` detector and the `beam_monitor_3` monitor."
13+
]
14+
},
15+
{
16+
"cell_type": "code",
17+
"execution_count": null,
18+
"id": "1",
19+
"metadata": {},
20+
"outputs": [],
21+
"source": [
22+
"import numpy as np\n",
23+
"import scipp as sc\n",
24+
"import scippnexus as sx\n",
25+
"import h5py as h5\n",
26+
"import shutil"
27+
]
28+
},
29+
{
30+
"cell_type": "code",
31+
"execution_count": null,
32+
"id": "2",
33+
"metadata": {},
34+
"outputs": [],
35+
"source": [
36+
"def load_mcstas_simulation_data(\n",
37+
" file_path,\n",
38+
" nevents,\n",
39+
" data_path=\"entry1/data/transmission_event_signal_dat_list_p_t_x_y_z_vx_vy_vz/events\",\n",
40+
"):\n",
41+
" nevents = int(nevents)\n",
42+
" with sx.File(file_path, \"r\") as f:\n",
43+
" # The name p_t_x_y_z_vx_vy_vz represents\n",
44+
" # probability, time of arrival, position(x, y, z) and velocity(vx, vy, vz).\n",
45+
" # The name also represents the order of each field in the table.\n",
46+
" # For example, probability is the first field, so data['dim_1', 0] is the probability.\n",
47+
" data = f[data_path][()].rename_dims({\"dim_0\": \"event\"})\n",
48+
"\n",
49+
" probabilities = data[\"dim_1\", 0].copy()\n",
50+
" probabilities.unit = \"dimensionless\"\n",
51+
" time_of_arrival = data[\"dim_1\", 1].copy()\n",
52+
" time_of_arrival.unit = \"s\" # Hardcoded unit from the data.\n",
53+
" positions = data[\"dim_1\", 2:5]\n",
54+
" # Units are hardcoded from the data.\n",
55+
" x_pos = positions[\"dim_1\", 0].copy()\n",
56+
" x_pos.unit = \"m\"\n",
57+
" y_pos = positions[\"dim_1\", 1].copy()\n",
58+
" y_pos.unit = \"m\"\n",
59+
"\n",
60+
" inds = np.random.choice(\n",
61+
" np.arange(data.sizes[\"event\"]),\n",
62+
" nevents,\n",
63+
" p=probabilities.values / probabilities.values.sum(),\n",
64+
" )\n",
65+
"\n",
66+
" out = sc.DataArray(\n",
67+
" data=sc.array(\n",
68+
" dims=probabilities.dims,\n",
69+
" values=np.ones(nevents),\n",
70+
" unit=\"counts\",\n",
71+
" ),\n",
72+
" coords={\n",
73+
" \"time_of_arrival\": sc.array(\n",
74+
" dims=time_of_arrival.dims,\n",
75+
" values=time_of_arrival.to(unit=\"us\").values[inds],\n",
76+
" unit=\"us\",\n",
77+
" ),\n",
78+
" \"x\": sc.array(\n",
79+
" dims=x_pos.dims,\n",
80+
" values=x_pos.values[inds],\n",
81+
" unit=x_pos.unit,\n",
82+
" ),\n",
83+
" \"y\": sc.array(\n",
84+
" dims=y_pos.dims,\n",
85+
" values=y_pos.values[inds],\n",
86+
" unit=y_pos.unit,\n",
87+
" ),\n",
88+
" },\n",
89+
" )\n",
90+
" return out"
91+
]
92+
},
93+
{
94+
"cell_type": "code",
95+
"execution_count": null,
96+
"id": "3",
97+
"metadata": {},
98+
"outputs": [],
99+
"source": [
100+
"def replace_dataset(entry, name, values):\n",
101+
" attrs = dict(entry[name].attrs)\n",
102+
" del entry[name]\n",
103+
" dset = entry.create_dataset(name, data=values)\n",
104+
" dset.attrs.update(attrs)\n",
105+
"\n",
106+
"\n",
107+
"def mcstas_to_nexus(\n",
108+
" mcstas_data_file: str,\n",
109+
" template_nexus_file: str,\n",
110+
" outfile: str,\n",
111+
" nevents: float = 1e6,\n",
112+
" detector_entry_path: str | None = \"entry/instrument/event_mode_detectors/timepix3\",\n",
113+
" monitor_entry_path: str | None = \"entry/instrument/beam_monitor_3\",\n",
114+
"):\n",
115+
" \"\"\"\n",
116+
" Store the events from a McStas Odin simulation in a NeXus CODA file.\n",
117+
"\n",
118+
" Parameters\n",
119+
" ----------\n",
120+
" mcstas_data_file:\n",
121+
" Data file containing simulated McStas events.\n",
122+
" template_nexus_file:\n",
123+
" NeXus file containing geometry and instrument info, used as a template.\n",
124+
" outfile:\n",
125+
" Output file to be written.\n",
126+
" nevents:\n",
127+
" Number of events to have in the output file\n",
128+
" (events are sampled from the probabilities of the mcstas events).\n",
129+
" detector_entry_path:\n",
130+
" Location of the event detector in the nexus tree.\n",
131+
" monitor_entry_path:\n",
132+
" Location of the event monitor in the nexus tree.\n",
133+
" \"\"\"\n",
134+
" # Find shape of detector panel\n",
135+
" with h5.File(template_nexus_file, \"r\") as f:\n",
136+
" shape = f[f\"{detector_entry_path}/x_pixel_offset\"].shape\n",
137+
" det_numbers = f[f\"{detector_entry_path}/detector_number\"][()]\n",
138+
"\n",
139+
" da = load_mcstas_simulation_data(mcstas_data_file, nevents=int(nevents))\n",
140+
" binned = da.bin(y=shape[0], x=shape[1]).rename_dims(y=\"dim_0\", x=\"dim_1\")\n",
141+
"\n",
142+
" toa = binned.bins.coords[\"time_of_arrival\"].bins.concat().value\n",
143+
"\n",
144+
" # IMPORTANT! we need to sort the arrays below according to toa,\n",
145+
" # so that the event_index does not get messed up!\n",
146+
" event_id = sc.sort(\n",
147+
" (\n",
148+
" sc.bins_like(binned, sc.array(dims=binned.dims, values=det_numbers))\n",
149+
" .bins.concat()\n",
150+
" .value\n",
151+
" ),\n",
152+
" key=toa,\n",
153+
" )\n",
154+
"\n",
155+
" unit = \"ns\"\n",
156+
" period = (1.0 / sc.scalar(14.0, unit=\"Hz\")).to(unit=unit)\n",
157+
" start = sc.datetime(\"2024-01-01T12:00:00.000000000\")\n",
158+
"\n",
159+
" event_time_zero = sc.sort(\n",
160+
" (period * (toa.to(unit=\"ns\", copy=False) // period)).to(dtype=int) + start,\n",
161+
" key=toa,\n",
162+
" )\n",
163+
"\n",
164+
" event_time_offset = sc.sort(toa % period.to(unit=toa.unit), key=toa)\n",
165+
"\n",
166+
" event_index = sc.DataArray(\n",
167+
" data=sc.ones_like(event_time_offset),\n",
168+
" coords={\"event_time_zero\": event_time_zero},\n",
169+
" ).group(\"event_time_zero\")\n",
170+
"\n",
171+
" event_index = sc.cumsum(event_index.bins.size())\n",
172+
" event_index.values = np.concatenate([[0], event_index.values[:-1]])\n",
173+
"\n",
174+
" # Now edit the template file\n",
175+
" print(f\"Writing {outfile} file\")\n",
176+
" shutil.copyfile(template_nexus_file, outfile)\n",
177+
" f = h5.File(outfile, \"r+\")\n",
178+
"\n",
179+
" # Detector data\n",
180+
" if detector_entry_path is not None:\n",
181+
" # remove translation offset\n",
182+
" f[f\"{detector_entry_path}/transformations/translation\"].attrs[\"offset\"] = (\n",
183+
" np.array([0, 0, 0], dtype=\"float32\")\n",
184+
" )\n",
185+
"\n",
186+
" event_data = f[f\"{detector_entry_path}/timepix3_events\"]\n",
187+
" replace_dataset(event_data, name=\"event_id\", values=event_id.values)\n",
188+
" replace_dataset(\n",
189+
" event_data,\n",
190+
" name=\"event_time_offset\",\n",
191+
" values=event_time_offset.to(\n",
192+
" unit=event_data[\"event_time_offset\"].attrs[\"units\"], copy=False\n",
193+
" ).values,\n",
194+
" )\n",
195+
" replace_dataset(event_data, name=\"event_index\", values=event_index.values)\n",
196+
" replace_dataset(\n",
197+
" event_data,\n",
198+
" name=\"event_time_zero\",\n",
199+
" values=event_index.coords[\"event_time_zero\"]\n",
200+
" .to(unit=event_data[\"event_time_zero\"].attrs[\"units\"], copy=False)\n",
201+
" .values.astype(int),\n",
202+
" )\n",
203+
"\n",
204+
" # Monitor data\n",
205+
" if monitor_entry_path is not None:\n",
206+
" monitor_data = f[f\"{monitor_entry_path}/monitor_3_events\"]\n",
207+
" replace_dataset(\n",
208+
" monitor_data, name=\"event_id\", values=np.zeros_like(event_id.values)\n",
209+
" )\n",
210+
" replace_dataset(\n",
211+
" monitor_data,\n",
212+
" name=\"event_time_offset\",\n",
213+
" values=event_time_offset.to(\n",
214+
" unit=monitor_data[\"event_time_offset\"].attrs[\"units\"], copy=False\n",
215+
" ).values,\n",
216+
" )\n",
217+
" replace_dataset(monitor_data, name=\"event_index\", values=event_index.values)\n",
218+
" replace_dataset(\n",
219+
" monitor_data,\n",
220+
" name=\"event_time_zero\",\n",
221+
" values=event_index.coords[\"event_time_zero\"]\n",
222+
" .to(unit=monitor_data[\"event_time_zero\"].attrs[\"units\"], copy=False)\n",
223+
" .values.astype(int),\n",
224+
" )"
225+
]
226+
},
227+
{
228+
"cell_type": "code",
229+
"execution_count": null,
230+
"id": "4",
231+
"metadata": {},
232+
"outputs": [],
233+
"source": [
234+
"files = [\n",
235+
" \"small_mcstas_sample_images.h5\",\n",
236+
" \"small_mcstas_ob_images.h5\",\n",
237+
" \"iron_simulation_sample.h5\",\n",
238+
" \"iron_simulation_ob.h5\",\n",
239+
"]\n",
240+
"\n",
241+
"for file in files:\n",
242+
" mcstas_to_nexus(\n",
243+
" mcstas_data_file=file,\n",
244+
" template_nexus_file=\"977695_00072982.hdf\",\n",
245+
" outfile=file.replace(\".h5\", \".nxs\"),\n",
246+
" nevents=1e6,\n",
247+
" )"
248+
]
249+
}
250+
],
251+
"metadata": {
252+
"kernelspec": {
253+
"display_name": "Python 3 (ipykernel)",
254+
"language": "python",
255+
"name": "python3"
256+
},
257+
"language_info": {
258+
"codemirror_mode": {
259+
"name": "ipython",
260+
"version": 3
261+
},
262+
"file_extension": ".py",
263+
"mimetype": "text/x-python",
264+
"name": "python",
265+
"nbconvert_exporter": "python",
266+
"pygments_lexer": "ipython3",
267+
"version": "3.12.7"
268+
}
269+
},
270+
"nbformat": 4,
271+
"nbformat_minor": 5
272+
}

0 commit comments

Comments
 (0)