eclipse 0.13.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (50) hide show
  1. eclipse/__init__.py +95 -0
  2. eclipse/analysis.py +853 -0
  3. eclipse/atmosphere.py +859 -0
  4. eclipse/band.py +477 -0
  5. eclipse/cli.py +138 -0
  6. eclipse/config.py +1710 -0
  7. eclipse/continuum.py +413 -0
  8. eclipse/data/eis_response/EIS_EffArea_A.004 +61 -0
  9. eclipse/data/eis_response/EIS_EffArea_B.004 +60 -0
  10. eclipse/data/eis_response/EIS_EffArea_B.005 +64 -0
  11. eclipse/data/eis_response/README_dz2025 +50 -0
  12. eclipse/data/eis_response/fit_eis_ea_2023-05-04_smooth.sav +0 -0
  13. eclipse/data/eis_response/source.txt +33 -0
  14. eclipse/data/science_cases.yaml +365 -0
  15. eclipse/data/throughput/grating_reflection_efficiency.dat +25 -0
  16. eclipse/data/throughput/primary_mirror_coating_reflectance.dat +25 -0
  17. eclipse/data/throughput/source.txt +6 -0
  18. eclipse/data/throughput/throughput_aluminium_1000_angstrom.dat +503 -0
  19. eclipse/data/throughput/throughput_aluminium_oxide_1000_angstrom.dat +503 -0
  20. eclipse/data/throughput/throughput_carbon_1000_angstrom.dat +503 -0
  21. eclipse/data_processing.py +904 -0
  22. eclipse/deprecated.py +161 -0
  23. eclipse/eis_calibration.py +573 -0
  24. eclipse/extern/__init__.py +0 -0
  25. eclipse/extern/mpfit.py +2388 -0
  26. eclipse/fitting.py +1734 -0
  27. eclipse/frame.py +594 -0
  28. eclipse/full_frame.py +944 -0
  29. eclipse/main.py +1646 -0
  30. eclipse/monte_carlo.py +660 -0
  31. eclipse/pinhole_diffraction.py +409 -0
  32. eclipse/psf.py +6 -0
  33. eclipse/radiometric.py +1328 -0
  34. eclipse/raster.py +869 -0
  35. eclipse/readout.py +1002 -0
  36. eclipse/resolution_budget.py +108 -0
  37. eclipse/results_file.py +1226 -0
  38. eclipse/sampling.py +351 -0
  39. eclipse/science_cases.py +364 -0
  40. eclipse/synthesis.py +2492 -0
  41. eclipse/synthesis_cli.py +12 -0
  42. eclipse/synthesis_file.py +882 -0
  43. eclipse/utils.py +1145 -0
  44. eclipse-0.13.0.dist-info/METADATA +70 -0
  45. eclipse-0.13.0.dist-info/RECORD +50 -0
  46. eclipse-0.13.0.dist-info/WHEEL +5 -0
  47. eclipse-0.13.0.dist-info/entry_points.txt +7 -0
  48. eclipse-0.13.0.dist-info/licenses/LICENSE +3 -0
  49. eclipse-0.13.0.dist-info/top_level.txt +2 -0
  50. euvst_response/__init__.py +129 -0
eclipse/synthesis.py ADDED
@@ -0,0 +1,2492 @@
1
+ import os
2
+ import sys
3
+ import re
4
+ import argparse
5
+ import warnings
6
+ from pathlib import Path
7
+ from typing import Dict, Tuple, List, Optional, Union
8
+ import numpy as np
9
+ from scipy.interpolate import RegularGridInterpolator
10
+ import astropy.units as u
11
+ import astropy.constants as const
12
+ from tqdm import tqdm
13
+ import psutil
14
+ import dill
15
+ from ndcube import NDCube
16
+ from astropy.wcs import WCS
17
+ from .utils import (angle_to_distance, element_data, require_uniform_grid,
18
+ require_downsample_divides, velocity_centers_to_edges, velocity_grid,
19
+ view_axis_and_side, view_name, OBSERVER_SIDE, VELOCITY_CONVENTION)
20
+ from .synthesis_file import write_line_cubes
21
+ from .continuum import compute_continuum_fiasco, continuum_spectra, continuum_windows
22
+ from .atmosphere import (AXES, NUMPY_AXIS, Atmosphere, _offer_database_build,
23
+ mass_per_electron, read_atmosphere, require_mass_per_electron)
24
+
25
+ ##############################################################################
26
+ # ---------------------------------------------------------------------------
27
+ # I/O helpers
28
+ # ---------------------------------------------------------------------------
29
+ ##############################################################################
30
+
31
+ def load_cube(
32
+ file_path: str | Path,
33
+ shape: Tuple[int, int, int] = (512, 768, 256),
34
+ unit: Optional[u.Unit] = None,
35
+ downsample: int | bool = False,
36
+ precision: type = np.float32,
37
+ voxel_dx: Optional[u.Quantity] = None,
38
+ voxel_dy: Optional[u.Quantity] = None,
39
+ voxel_dz: Optional[u.Quantity] = None,
40
+ create_ndcube: bool = False,
41
+ ) -> np.ndarray | u.Quantity | NDCube:
42
+ """
43
+ Read a Fortran-ordered float32 binary cube, optionally as an NDCube.
44
+
45
+ The file stores the cube as (x, z, y). It is returned as (z, y, x), so
46
+ `data[k]` is a horizontal slice indexed `[y, x]`.
47
+
48
+ Parameters
49
+ ----------
50
+ file_path : str | Path
51
+ Path to the binary file.
52
+ shape : Tuple[int, int, int]
53
+ Full cube dimensions in the file's storage order, `(nx, nz, ny)`:
54
+ the vertical axis is second, not last.
55
+ unit : astropy.units.Unit, optional
56
+ Unit to attach (e.g. u.K or u.g/u.cm**3). If None, returns a plain
57
+ ndarray.
58
+ downsample : int | bool
59
+ Integer factor; if not False, keep every `downsample`-th cell along
60
+ each axis (simple stride).
61
+ precision : type
62
+ np.float32 or np.float64 for the returned dtype.
63
+ voxel_dx, voxel_dy, voxel_dz : u.Quantity, optional
64
+ Voxel sizes in the file, at full resolution. With `downsample`, the
65
+ returned cube's WCS uses them multiplied by it. Required if
66
+ `create_ndcube` is True.
67
+ create_ndcube : bool, optional
68
+ If True, return an NDCube with WCS coordinates.
69
+
70
+ Returns
71
+ -------
72
+ ndarray, Quantity, or NDCube
73
+ Array of shape (nz', ny', nx'), or an NDCube with its coordinates.
74
+ """
75
+ if downsample:
76
+ require_downsample_divides(shape, downsample)
77
+
78
+ data = np.fromfile(file_path, dtype=np.float32).reshape(shape, order="F")
79
+ data = data.transpose(1, 2, 0) # (z,y,x)
80
+
81
+ if downsample:
82
+ data = data[::downsample, ::downsample, ::downsample]
83
+ # Each kept cell spans `downsample` cells of the file. Make new
84
+ # quantities rather than use *=, which would scale the caller's voxel
85
+ # sizes and compound across calls.
86
+ voxel_dx, voxel_dy, voxel_dz = (
87
+ None if size is None else size * downsample
88
+ for size in (voxel_dx, voxel_dy, voxel_dz))
89
+
90
+ data = data.astype(precision, copy=False)
91
+
92
+ if unit is not None:
93
+ data = data * unit
94
+
95
+ if create_ndcube:
96
+ return create_atmosphere_ndcube(data, voxel_dx, voxel_dy, voxel_dz)
97
+ else:
98
+ return data
99
+
100
+
101
+ def create_atmosphere_ndcube(
102
+ data: np.ndarray | u.Quantity,
103
+ voxel_dx: u.Quantity,
104
+ voxel_dy: u.Quantity,
105
+ voxel_dz: u.Quantity,
106
+ ) -> NDCube:
107
+ """
108
+ NDCube of a simulation cube, with solar coordinates from its cell sizes.
109
+
110
+ x and y are centred on zero, and z starts at zero.
111
+
112
+ Parameters
113
+ ----------
114
+ data : np.ndarray or u.Quantity
115
+ Cube of shape (nz, ny, nx); `data[k]` is a horizontal slice indexed
116
+ `[y, x]`.
117
+ voxel_dx, voxel_dy, voxel_dz : u.Quantity
118
+ Cell size along each axis.
119
+
120
+ Returns
121
+ -------
122
+ NDCube
123
+ The cube, with coordinates in Mm.
124
+ """
125
+ nz, ny, nx = data.shape
126
+
127
+ # Create WCS for heliocentric coordinates
128
+ wcs = WCS(naxis=3)
129
+ wcs.wcs.ctype = ['SOLX', 'SOLY', 'SOLZ']
130
+ wcs.wcs.cunit = ['Mm', 'Mm', 'Mm']
131
+
132
+ # Reference pixels (1-indexed for WCS)
133
+ wcs.wcs.crpix = [(nx + 1) / 2, (ny + 1) / 2, 1] # Z starts at first pixel
134
+
135
+ # Reference values
136
+ wcs.wcs.crval = [0, 0, 0] # X,Y centered at origin, Z starts at 0
137
+
138
+ # Pixel scales
139
+ wcs.wcs.cdelt = [
140
+ voxel_dx.to(u.Mm).value,
141
+ voxel_dy.to(u.Mm).value,
142
+ voxel_dz.to(u.Mm).value
143
+ ]
144
+
145
+ # A plain array has no unit.
146
+ return NDCube(np.asarray(getattr(data, "value", data)),
147
+ wcs=wcs,
148
+ unit=getattr(data, "unit", None))
149
+
150
+
151
+ def read_timestep_time(file_path: Path) -> float:
152
+ """
153
+ Read simulation time from MHD header file.
154
+
155
+ The header file contains a single line with 9 space-separated values:
156
+ nx ny nz dx dy dz time dt va_max
157
+
158
+ The time is the 7th value (index 6) in seconds.
159
+
160
+ Parameters
161
+ ----------
162
+ file_path : Path
163
+ Path to the header file.
164
+
165
+ Returns
166
+ -------
167
+ float
168
+ Simulation time in seconds.
169
+ """
170
+ with open(file_path, 'r') as f:
171
+ line = f.read().strip()
172
+ values = line.split()
173
+ if len(values) < 7:
174
+ raise ValueError(f"Header file has insufficient values: {file_path} ({len(values)} values)")
175
+ return float(values[6])
176
+
177
+
178
+ def discover_timesteps(
179
+ time_dir: Path,
180
+ time_filename: str,
181
+ ) -> Dict[str, float]:
182
+ """
183
+ Discover all available timesteps and their simulation times.
184
+
185
+ Parameters
186
+ ----------
187
+ time_dir : Path
188
+ Directory containing header files.
189
+ time_filename : str
190
+ Filename prefix before the timestep suffix (e.g., "Header").
191
+
192
+ Returns
193
+ -------
194
+ dict
195
+ Mapping of timestep suffix to simulation time in seconds.
196
+ E.g., {"0270000": 26729.535, "0280000": 27571.395, ...}
197
+ """
198
+ if not time_dir.is_dir():
199
+ raise FileNotFoundError(f"Time directory not found: {time_dir}")
200
+
201
+ timesteps = {}
202
+
203
+ for file_path in sorted(time_dir.iterdir()):
204
+ if file_path.is_file() and file_path.name.startswith(time_filename):
205
+ # Extract suffix after the filename prefix
206
+ suffix = file_path.name[len(time_filename):]
207
+ if suffix.startswith('.'):
208
+ suffix = suffix[1:] # Remove leading dot
209
+
210
+ try:
211
+ sim_time = read_timestep_time(file_path)
212
+ timesteps[suffix] = sim_time
213
+ except Exception as e:
214
+ warnings.warn(f"Could not read time from {file_path}: {e}")
215
+
216
+ if not timesteps:
217
+ raise ValueError(f"No valid timestep files found in {time_dir} with prefix '{time_filename}'")
218
+
219
+ return timesteps
220
+
221
+
222
+ def get_file_for_timestep(
223
+ directory: Path,
224
+ filename: str,
225
+ suffix: str,
226
+ ) -> Path:
227
+ """
228
+ Get the file path for a specific timestep.
229
+
230
+ Parameters
231
+ ----------
232
+ directory : Path
233
+ Directory containing the data files.
234
+ filename : str
235
+ Filename prefix before the suffix (e.g., "eosT").
236
+ suffix : str
237
+ Timestep suffix (e.g., "0270000").
238
+
239
+ Returns
240
+ -------
241
+ Path
242
+ Full path to the file.
243
+ """
244
+ file_path = directory / f"{filename}.{suffix}"
245
+ if not file_path.exists():
246
+ raise FileNotFoundError(f"File not found: {file_path}")
247
+ return file_path
248
+
249
+
250
+ def compute_slice_timestep_mapping_mhd(
251
+ nx_mhd: int,
252
+ voxel_dx: u.Quantity,
253
+ slit_width: u.Quantity,
254
+ slit_rest_time: u.Quantity,
255
+ timestep_times: Dict[str, float],
256
+ crop_x: Optional[Tuple[u.Quantity, u.Quantity]] = None,
257
+ ) -> Tuple[List[str], Dict[str, List[int]]]:
258
+ """
259
+ Timestep suffix to use for each MHD X-slice.
260
+
261
+ The spectrometer scans from right to left (high X to low X), and each
262
+ slit position covers `slit_width` converted to Mm. Each MHD slice is
263
+ assigned to a slit position, then to the timestep at that position's
264
+ observation time. The time is worked out from the slice's absolute
265
+ position, so errors do not accumulate.
266
+
267
+ With `crop_x`, the scan starts at the right edge of the crop region.
268
+ Slices to the right of it get the first timestep.
269
+
270
+ Parameters
271
+ ----------
272
+ nx_mhd : int
273
+ Number of MHD X-slices.
274
+ voxel_dx : u.Quantity
275
+ MHD voxel size along X, as a length (e.g. Mm).
276
+ slit_width : u.Quantity
277
+ Slit width as an angle (e.g. arcsec).
278
+ slit_rest_time : u.Quantity
279
+ Slit rest time per position.
280
+ timestep_times : dict
281
+ Simulation time in seconds for each timestep suffix.
282
+ crop_x : tuple of u.Quantity, optional
283
+ (x_min, x_max) crop boundaries. The scan starts at x_max, the right
284
+ edge of the crop; slices to the right of it use the first timestep.
285
+
286
+ Returns
287
+ -------
288
+ slice_mapping : list
289
+ Timestep suffix for each MHD X-slice.
290
+ grouped : dict
291
+ MHD slice indices for each timestep suffix.
292
+ """
293
+ rest_time_sec = slit_rest_time.to_value(u.s)
294
+
295
+ # Convert slit width from angular to physical distance
296
+ slit_physical = angle_to_distance(slit_width).to(u.Mm)
297
+ voxel_physical = voxel_dx.to(u.Mm)
298
+
299
+ # Sort timesteps by simulation time
300
+ sorted_timesteps = sorted(timestep_times.items(), key=lambda x: x[1])
301
+ suffixes = [s for s, t in sorted_timesteps]
302
+ times = np.array([t for s, t in sorted_timesteps])
303
+ first_suffix = suffixes[0]
304
+
305
+ # Check if observation time range extends beyond available MHD time range
306
+ t0 = times[0]
307
+ t_final = times[-1]
308
+ mhd_duration = t_final - t0
309
+
310
+ # Total physical extent of the full domain
311
+ total_extent = nx_mhd * voxel_physical
312
+
313
+ # The scan starts at the right edge of the domain, or of the crop if given.
314
+ if crop_x is not None:
315
+ x_min_crop = crop_x[0].to(u.Mm)
316
+ x_max_crop = crop_x[1].to(u.Mm)
317
+
318
+ # The WCS centres X on zero (x=0 is at nx/2), so convert to absolute
319
+ # positions from the left edge.
320
+ domain_center = total_extent / 2
321
+ scan_start_position = domain_center + x_max_crop # Right edge of crop in absolute coords
322
+ scan_end_position = domain_center + x_min_crop # Left edge of crop in absolute coords
323
+
324
+ # Calculate observation duration for just the cropped region
325
+ crop_extent = x_max_crop - x_min_crop
326
+ n_slit_positions_crop = int(np.ceil((crop_extent / slit_physical).decompose().value))
327
+ observation_duration = n_slit_positions_crop * rest_time_sec
328
+
329
+ print(f" Crop region: X = [{x_min_crop:.3f}, {x_max_crop:.3f}]")
330
+ print(f" Observation starts at X = {x_max_crop:.3f} (right edge of crop)")
331
+ else:
332
+ scan_start_position = total_extent # Right edge of full domain
333
+ scan_end_position = 0 * u.Mm
334
+ n_slit_positions_crop = int(np.ceil((total_extent / slit_physical).decompose().value))
335
+ observation_duration = n_slit_positions_crop * rest_time_sec
336
+
337
+ if observation_duration > mhd_duration:
338
+ print(f" WARNING: Observation duration ({observation_duration:.1f} s) exceeds MHD time range ({mhd_duration:.1f} s)")
339
+ print(f" Slices observed after t={t_final:.1f} s will use the last available timestep")
340
+
341
+ slice_mapping = []
342
+
343
+ # Find the slit position of each MHD X-slice. The scan runs right to left,
344
+ # so scan_start_position is observed first, at t=t0.
345
+ for mhd_slice_idx in range(nx_mhd):
346
+ # Physical position of this slice (center of voxel) in absolute coords
347
+ x_physical = (mhd_slice_idx + 0.5) * voxel_physical
348
+
349
+ # Distance from the scan start position (right edge of observation region)
350
+ distance_from_scan_start = scan_start_position - x_physical
351
+
352
+ if distance_from_scan_start < 0:
353
+ # Outside the crop on the right: use the first timestep.
354
+ slice_mapping.append(first_suffix)
355
+ continue
356
+
357
+ # Slit position 0 is at scan_start_position.
358
+ slit_position = int(np.floor((distance_from_scan_start / slit_physical).decompose().value))
359
+
360
+ # Observation time for this slit position
361
+ observation_time = t0 + slit_position * rest_time_sec
362
+
363
+ # Find the latest timestep that doesn't exceed observation_time
364
+ idx = np.searchsorted(times, observation_time, side='right') - 1
365
+ # Clamp to valid range (0 to len(times)-1)
366
+ idx = max(0, min(idx, len(times) - 1))
367
+ slice_mapping.append(suffixes[idx])
368
+
369
+ # Group slices by timestep for efficient processing
370
+ grouped = {}
371
+ for slice_idx, suffix in enumerate(slice_mapping):
372
+ if suffix not in grouped:
373
+ grouped[suffix] = []
374
+ grouped[suffix].append(slice_idx)
375
+
376
+ # Print statistics
377
+ if crop_x is not None:
378
+ print(f" Cropped region extent: {crop_extent:.3f}")
379
+ else:
380
+ print(f" Physical domain extent: {total_extent:.3f}")
381
+ print(f" Slit physical width: {slit_physical:.3f}")
382
+ print(f" Number of slit positions: {n_slit_positions_crop}")
383
+ print(f" Total raster time: {observation_duration:.1f} s")
384
+
385
+ return slice_mapping, grouped
386
+
387
+
388
+ def apply_cube_cropping(
389
+ temp_cube: NDCube,
390
+ rho_cube: NDCube,
391
+ vel_cube: NDCube,
392
+ crop_x: Optional[List[str]],
393
+ crop_y: Optional[List[str]],
394
+ crop_z: Optional[List[str]],
395
+ ) -> Tuple[NDCube, NDCube, NDCube]:
396
+ """
397
+ Apply cropping to temperature, density, and velocity cubes.
398
+
399
+ Parameters
400
+ ----------
401
+ temp_cube, rho_cube, vel_cube : NDCube
402
+ Input cubes to crop.
403
+ crop_x, crop_y, crop_z : list of str or None
404
+ Crop boundaries for each axis as [min, max] strings with units.
405
+
406
+ Returns
407
+ -------
408
+ tuple of NDCube
409
+ Cropped (temp_cube, rho_cube, vel_cube).
410
+ """
411
+ # Crop points are given in world axis order, which is (SOLX, SOLY, SOLZ).
412
+ point1 = []
413
+ point2 = []
414
+
415
+ if crop_x:
416
+ point1.append(u.Quantity(crop_x[0]))
417
+ point2.append(u.Quantity(crop_x[1]))
418
+ else:
419
+ point1.append(None)
420
+ point2.append(None)
421
+
422
+ if crop_y:
423
+ point1.append(u.Quantity(crop_y[0]))
424
+ point2.append(u.Quantity(crop_y[1]))
425
+ else:
426
+ point1.append(None)
427
+ point2.append(None)
428
+
429
+ if crop_z:
430
+ point1.append(u.Quantity(crop_z[0]))
431
+ point2.append(u.Quantity(crop_z[1]))
432
+ else:
433
+ point1.append(None)
434
+ point2.append(None)
435
+
436
+ # keepdims keeps an axis cropped to one cell, since the synthesis needs
437
+ # all three axes.
438
+ temp_cube = temp_cube.crop(point1, point2, keepdims=True)
439
+ rho_cube = rho_cube.crop(point1, point2, keepdims=True)
440
+ vel_cube = vel_cube.crop(point1, point2, keepdims=True)
441
+
442
+ return temp_cube, rho_cube, vel_cube
443
+
444
+
445
+ def build_composite_cubes_mhd(
446
+ base_dir: Path,
447
+ temp_dir: str,
448
+ temp_filename: str,
449
+ rho_dir: str,
450
+ rho_filename: str,
451
+ vel_dir: str,
452
+ vel_filename: str,
453
+ slice_mapping: List[str],
454
+ grouped_slices: Dict[str, List[int]],
455
+ cube_shape: Tuple[int, int, int],
456
+ voxel_dx: u.Quantity,
457
+ voxel_dy: u.Quantity,
458
+ voxel_dz: u.Quantity,
459
+ downsample: int | bool,
460
+ precision: type,
461
+ ) -> Tuple[NDCube, NDCube, NDCube]:
462
+ """
463
+ Build composite temp, rho, vel cubes from multiple timesteps at MHD resolution.
464
+
465
+ Each MHD X-slice comes from its timestep in `slice_mapping`. There is no
466
+ spatial rebinning: the output is at full MHD resolution.
467
+
468
+ Parameters
469
+ ----------
470
+ base_dir : Path
471
+ Base directory for atmosphere data.
472
+ temp_dir, temp_filename : str
473
+ Directory and filename prefix for temperature files.
474
+ rho_dir, rho_filename : str
475
+ Directory and filename prefix for density files.
476
+ vel_dir, vel_filename : str
477
+ Directory and filename prefix for velocity files.
478
+ slice_mapping : list
479
+ Timestep suffix for each MHD X-slice.
480
+ grouped_slices : dict
481
+ Slices grouped by timestep for efficient processing.
482
+ cube_shape : tuple
483
+ Original cube dimensions in the file's storage order.
484
+ voxel_dx, voxel_dy, voxel_dz : u.Quantity
485
+ Voxel sizes in the files, at full resolution; `load_cube` applies the
486
+ downsampling.
487
+ downsample : int or bool
488
+ Downsampling factor.
489
+ precision : type
490
+ Numerical precision.
491
+
492
+ Returns
493
+ -------
494
+ tuple
495
+ (temp_composite, rho_composite, vel_composite) NDCubes at MHD resolution.
496
+ """
497
+ nx_mhd = len(slice_mapping)
498
+
499
+ # Load first timestep to determine dimensions and get reference WCS
500
+ first_suffix = list(grouped_slices.keys())[0]
501
+ temp_file = get_file_for_timestep(base_dir / temp_dir, temp_filename, first_suffix)
502
+ temp_cube_ref = load_cube(
503
+ temp_file, shape=cube_shape, unit=u.K,
504
+ downsample=downsample, precision=precision,
505
+ voxel_dx=voxel_dx, voxel_dy=voxel_dy, voxel_dz=voxel_dz,
506
+ create_ndcube=True
507
+ )
508
+
509
+ nz, ny, nx = temp_cube_ref.data.shape
510
+ reference_wcs = temp_cube_ref.wcs
511
+
512
+ # Verify dimensions match slice mapping
513
+ if nx != nx_mhd:
514
+ raise ValueError(f"Cube X dimension ({nx}) doesn't match slice mapping ({nx_mhd})")
515
+
516
+ # Initialise composite arrays
517
+ temp_composite = np.zeros((nz, ny, nx), dtype=precision)
518
+ rho_composite = np.zeros((nz, ny, nx), dtype=precision)
519
+ vel_composite = np.zeros((nz, ny, nx), dtype=precision)
520
+
521
+ # Process each timestep
522
+ for suffix, slice_indices in tqdm(grouped_slices.items(), desc="Loading timesteps", unit="timestep"):
523
+ # Load cubes for this timestep
524
+ temp_file = get_file_for_timestep(base_dir / temp_dir, temp_filename, suffix)
525
+ rho_file = get_file_for_timestep(base_dir / rho_dir, rho_filename, suffix)
526
+ vel_file = get_file_for_timestep(base_dir / vel_dir, vel_filename, suffix)
527
+
528
+ temp_cube = load_cube(
529
+ temp_file, shape=cube_shape, unit=u.K,
530
+ downsample=downsample, precision=precision,
531
+ voxel_dx=voxel_dx, voxel_dy=voxel_dy, voxel_dz=voxel_dz,
532
+ create_ndcube=True
533
+ )
534
+ rho_cube = load_cube(
535
+ rho_file, shape=cube_shape, unit=u.g/u.cm**3,
536
+ downsample=downsample, precision=precision,
537
+ voxel_dx=voxel_dx, voxel_dy=voxel_dy, voxel_dz=voxel_dz,
538
+ create_ndcube=True
539
+ )
540
+ vel_cube = load_cube(
541
+ vel_file, shape=cube_shape, unit=u.cm/u.s,
542
+ downsample=downsample, precision=precision,
543
+ voxel_dx=voxel_dx, voxel_dy=voxel_dy, voxel_dz=voxel_dz,
544
+ create_ndcube=True
545
+ )
546
+
547
+ # Copy relevant slices to composite (no rebinning - direct copy)
548
+ for slice_idx in slice_indices:
549
+ temp_composite[:, :, slice_idx] = temp_cube.data[:, :, slice_idx]
550
+ rho_composite[:, :, slice_idx] = rho_cube.data[:, :, slice_idx]
551
+ vel_composite[:, :, slice_idx] = vel_cube.data[:, :, slice_idx]
552
+
553
+ # Create NDCubes with proper WCS (at MHD resolution)
554
+ temp_ndcube = NDCube(temp_composite * u.K, wcs=reference_wcs, meta={"source": "composite_dynamic"})
555
+ rho_ndcube = NDCube(rho_composite * (u.g/u.cm**3), wcs=reference_wcs, meta={"source": "composite_dynamic"})
556
+ vel_ndcube = NDCube(vel_composite * (u.cm/u.s), wcs=reference_wcs, meta={"source": "composite_dynamic"})
557
+
558
+ return temp_ndcube, rho_ndcube, vel_ndcube
559
+
560
+
561
+ def _match_candidates(line_name: str, wavelengths_aa: np.ndarray, observed: np.ndarray) -> tuple:
562
+ """
563
+ Indices of the transitions matching `line_name`, and how they were matched.
564
+
565
+ The match is, in order of preference:
566
+
567
+ - Transitions whose wavelength, rounded to the decimals in the name,
568
+ equals the name's: observed ones if any, else theoretical ones, so a
569
+ line with only a theoretical wavelength can still be named.
570
+ - The nearest observed line. CHIANTI's theoretical wavelengths include
571
+ many weak transitions within a few mA of strong lines, and a name
572
+ slightly off an observed wavelength (as from another line list) would
573
+ otherwise pick one of those, up to 1e12 times fainter.
574
+ - The nearest theoretical line, for an ion with no observed lines.
575
+
576
+ All transitions at the chosen wavelength are returned, so the brightest
577
+ can be taken.
578
+ """
579
+ number = line_name.split("_", 1)[1]
580
+ decimals = len(number.partition(".")[2])
581
+ target = float(number)
582
+ named = np.round(wavelengths_aa, decimals) == np.round(target, decimals)
583
+ if (named & observed).any():
584
+ pool, how = named & observed, "named"
585
+ elif named.any():
586
+ pool, how = named, "named"
587
+ elif observed.any():
588
+ pool, how = observed, "nearest observed"
589
+ else:
590
+ pool, how = np.ones_like(named), "nearest theoretical"
591
+ distance = np.where(pool, np.abs(wavelengths_aa - target), np.inf)
592
+ return np.flatnonzero(distance == distance.min()), how
593
+
594
+
595
+ def _compute_single_ion(args):
596
+ """Worker that computes G(T,N) for one ion.
597
+
598
+ fiasco is imported here so that each spawned process gets its own HDF5
599
+ handles. With a temperature chunk, fiasco is called on that many
600
+ temperatures at a time and only the requested lines are kept, so memory
601
+ use is that of one chunk rather than the whole temperature grid."""
602
+
603
+ import fiasco
604
+ import logging
605
+
606
+ (elem, stage, temperature_K, densities_cm3, abundance, lines,
607
+ hdf5_dbase_root, temperature_chunk) = args
608
+ densities = densities_cm3 / u.cm**3
609
+ n_temperatures = len(temperature_K)
610
+ step = n_temperatures if temperature_chunk is None else temperature_chunk
611
+
612
+ # A spawned process re-imports fiasco and re-reads ~/.fiasco/fiascorc,
613
+ # so it does not see a database the parent selected in memory. The
614
+ # database root is passed in the arguments instead.
615
+ ion_kwargs = {} if hdf5_dbase_root is None else {
616
+ 'hdf5_dbase_root': hdf5_dbase_root}
617
+
618
+ # Silence fiasco's repeated warnings about missing proton data and
619
+ # autoionization/rrlvl files. These are gaps in CHIANTI (not all ions
620
+ # have .psplups or .auto/.rrlvl data), and fiasco and IDL both fall back
621
+ # to leaving out proton rates or using the single-ion model. The warnings
622
+ # fire once per density point, giving hundreds of identical lines.
623
+ fiasco_logger = logging.getLogger('fiasco')
624
+ prev_level = fiasco_logger.level
625
+ fiasco_logger.setLevel(logging.ERROR)
626
+
627
+ g_parts = {}
628
+ try:
629
+ for start in range(0, n_temperatures, step):
630
+ ion = fiasco.Ion(f'{elem} {stage}',
631
+ temperature_K[start:start + step] * u.K,
632
+ abundance=abundance, **ion_kwargs)
633
+
634
+ g = ion.contribution_function(densities)
635
+ pe_ratio = ion.proton_electron_ratio
636
+ g = g * pe_ratio[:, np.newaxis, np.newaxis]
637
+
638
+ # The transitions do not depend on temperature, so the lines are
639
+ # matched once, on the first chunk.
640
+ if start == 0:
641
+ bound = ion.transitions.is_bound_bound
642
+ bb_wl = ion.transitions.wavelength[bound]
643
+ # CHIANTI stores theoretical wavelengths as negative numbers;
644
+ # fiasco gives them positive and marks observed ones with this
645
+ # flag.
646
+ observed = np.asarray(ion.transitions.is_observed[bound])
647
+ matches = {line_name: _match_candidates(line_name, bb_wl.to_value(u.AA), observed)
648
+ for line_name, _ in lines}
649
+ g_parts = {(line_name, int(idx)): [] for line_name, (indices, _) in matches.items()
650
+ for idx in indices}
651
+ for (line_name, idx), parts in g_parts.items():
652
+ parts.append(g[:, :, idx].to(u.erg * u.cm**3 / u.s).value)
653
+ del g
654
+ finally:
655
+ fiasco_logger.setLevel(prev_level)
656
+
657
+ results = {}
658
+ for line_name, target_wl_aa in lines:
659
+ target_wl = target_wl_aa * u.AA
660
+ indices, how = matches[line_name]
661
+ tables = {}
662
+ for idx in indices:
663
+ g_tn = np.concatenate(g_parts[(line_name, int(idx))]).T
664
+ tables[int(idx)] = np.nan_to_num(g_tn, nan=0.0, posinf=0.0, neginf=0.0)
665
+ # Of the transitions at one wavelength, take the brightest.
666
+ idx = max(tables, key=lambda i: tables[i].max())
667
+ matched_wl = bb_wl[idx]
668
+
669
+ results[line_name] = {
670
+ "g_tn": tables[idx],
671
+ "atom": int(ion.atomic_number),
672
+ "ion": stage,
673
+ "target_wl_cm": float(target_wl.to(u.cm).value),
674
+ "matched_wl_aa": float(matched_wl.to(u.AA).value),
675
+ "match": how,
676
+ "observed": bool(observed[idx]),
677
+ "transition": idx,
678
+ "delta_aa": float(abs(matched_wl - target_wl).to(u.AA).value),
679
+ # Database root this Ion was built with. fiasco sets it to the
680
+ # fiascorc default if none was chosen, so this is always the
681
+ # database the contribution functions came from.
682
+ "hdf5_dbase_root": str(ion.hdf5_dbase_root),
683
+ }
684
+ return results
685
+
686
+
687
+ # Line name: element, ionisation stage and wavelength in Angstrom, such as
688
+ # Fe12_195.1190.
689
+ _LINE_NAME = re.compile(r'^([A-Z][a-z]?)(\d+)_(\d+\.?\d*)$')
690
+
691
+
692
+ def _parse_line_name(name: str) -> Tuple[str, int, float]:
693
+ """
694
+ Element, ionisation stage and wavelength in Angstrom of a line name.
695
+
696
+ Names look like `Fe12_195.1190`. The element must exist, the stage must
697
+ be from 1 (neutral) to its atomic number plus one, and the wavelength
698
+ must be above zero. These are checked here because a bad element or
699
+ stage would otherwise only be found when the ion is made, after the
700
+ atmosphere has been read, and a wavelength of zero would match the
701
+ shortest transition.
702
+ """
703
+ match = _LINE_NAME.match(name)
704
+ if not match:
705
+ raise ValueError(f"Cannot parse line name '{name}'. Expected format like "
706
+ f"'Fe12_195.1190'.")
707
+ elem, stage, wavelength = match.group(1), int(match.group(2)), float(match.group(3))
708
+ try:
709
+ atomic_number, _ = element_data(elem)
710
+ except ValueError:
711
+ raise ValueError(f"Line name '{name}': {elem} is not an element.") from None
712
+ if not 1 <= stage <= atomic_number + 1:
713
+ raise ValueError(f"Line name '{name}': {elem} has ionisation stages 1 to "
714
+ f"{atomic_number + 1}, got {stage}.")
715
+ if not (np.isfinite(wavelength) and wavelength > 0):
716
+ raise ValueError(f"Line name '{name}': the wavelength must be above zero, in "
717
+ f"Angstrom.")
718
+ return elem, stage, wavelength
719
+
720
+ # Default grids for G(T, n_e): log10 T from 4 to 9 in steps of 0.05, and
721
+ # log10 n_e from 7 to 13 in steps of 0.3. density_grid extends the density
722
+ # grid on the same points as far as an atmosphere needs.
723
+ _LOGT_MIN, _LOGT_MAX, _N_T = 4.0, 9.0, 101
724
+ _LOGN_MIN, _LOGN_MAX, _N_N = 7.0, 13.0, 21
725
+
726
+
727
+ def compute_goft_fiasco(
728
+ line_names: List[str],
729
+ abundance: str = "sun_coronal_2021_chianti",
730
+ logT_min: float = _LOGT_MIN,
731
+ logT_max: float = _LOGT_MAX,
732
+ nT: int = _N_T,
733
+ logN_min: float = _LOGN_MIN,
734
+ logN_max: float = _LOGN_MAX,
735
+ nN: int = _N_N,
736
+ precision: type = np.float64,
737
+ n_workers: int = 0,
738
+ hdf5_dbase_root=None,
739
+ temperature_chunk: int | None = None,
740
+ ) -> Tuple[Dict[str, dict], np.ndarray, np.ndarray]:
741
+ """
742
+ Contribution function G(T, n_e) of each line, from CHIANTI via fiasco.
743
+
744
+ Each line is matched to a CHIANTI transition as described in the docs
745
+ page on naming spectral lines, and placed at CHIANTI's wavelength, `wl0`,
746
+ not the name's. G includes the proton-to-electron ratio, so it multiplies
747
+ an emission measure of n_e^2 dh. With two or more ions and more than one
748
+ worker, the ions are computed in separate processes, so a calling script
749
+ needs an `if __name__ == "__main__":` guard.
750
+
751
+ Parameters
752
+ ----------
753
+ line_names : list of str
754
+ Line names, such as `["Fe12_195.1190", "Fe09_171.073"]`.
755
+ abundance : str, optional
756
+ CHIANTI abundance set. Default `"sun_coronal_2021_chianti"`.
757
+ logT_min, logT_max : float, optional
758
+ Range of log10(T / K).
759
+ nT : int, optional
760
+ Number of temperatures, evenly spaced in log10(T) across that range.
761
+ logN_min, logN_max : float, optional
762
+ Range of log10(n_e / cm-3).
763
+ nN : int, optional
764
+ Number of densities, evenly spaced in log10(n_e) across that range.
765
+ precision : type, optional
766
+ `np.float32` or `np.float64` (default).
767
+ n_workers : int, optional
768
+ Number of worker processes, each computing one ion at a time, capped
769
+ at the number of ions. Default 0 uses every CPU this process may use
770
+ (as allocated by SLURM), up to that cap.
771
+ hdf5_dbase_root : str or Path, optional
772
+ CHIANTI database to read. Default is fiasco's, set in
773
+ `~/.fiasco/fiascorc`.
774
+ temperature_chunk : int, optional
775
+ Pass fiasco this many temperatures at a time, rather than the whole
776
+ grid, to use less memory. The result is the same. Default None (the
777
+ whole grid).
778
+
779
+ Returns
780
+ -------
781
+ goft_dict : dict
782
+ Entry per line name, holding `wl0` (CHIANTI wavelength), `g_tn` (G on
783
+ the density and temperature grid, shape (nN, nT), in erg cm3 / s),
784
+ `atom` and `ion` (atomic number and ionisation stage) and
785
+ `hdf5_dbase_root` (database used).
786
+ logT_grid : np.ndarray
787
+ Temperatures, as log10(T / K).
788
+ logN_grid : np.ndarray
789
+ Densities, as log10(n_e / cm-3).
790
+ """
791
+ # fiasco's contribution_function, G = Ab(X) f_{X,k} (N_j / N) A_ij dE_ij
792
+ # / n_e, leaves out n_H / n_e, which is multiplied in here.
793
+ # Workers are spawned rather than forked, for HDF5 safety, so each one
794
+ # imports fiasco and reads fiascorc afresh and does not see
795
+ # fiasco.defaults set in the parent. The database is therefore passed
796
+ # down, and each worker reports the root its Ion used, which is checked
797
+ # against the request.
798
+ # fiasco solves the level populations for all its temperatures at once,
799
+ # so its memory grows with their number. Temperature chunks are computed
800
+ # one after another, at the cost of reading the atomic data again.
801
+ if temperature_chunk is not None and temperature_chunk < 1:
802
+ raise ValueError(
803
+ f"temperature_chunk must be a positive number of temperatures, "
804
+ f"not {temperature_chunk}."
805
+ )
806
+
807
+ logT_grid = np.linspace(logT_min, logT_max, nT)
808
+ logN_grid = np.linspace(logN_min, logN_max, nN)
809
+
810
+ temperature_K = 10.0 ** logT_grid
811
+ densities_cm3 = 10.0 ** logN_grid
812
+
813
+ # ---- parse line names and group by ion for efficiency ----
814
+ ion_lines: Dict[Tuple[str, int], List[Tuple[str, float]]] = {}
815
+
816
+ for name in line_names:
817
+ elem, stage, wl = _parse_line_name(name)
818
+ ion_lines.setdefault((elem, stage), []).append((name, wl))
819
+
820
+ # Build worker arguments (all picklable plain types / numpy arrays)
821
+ dbase_root = None if hdf5_dbase_root is None else str(hdf5_dbase_root)
822
+ worker_args = [
823
+ (elem, stage, temperature_K, densities_cm3, abundance, lines, dbase_root,
824
+ temperature_chunk)
825
+ for (elem, stage), lines in ion_lines.items()
826
+ ]
827
+ if temperature_chunk is not None and temperature_chunk < nT:
828
+ print(f" {nT} temperatures in chunks of {temperature_chunk}")
829
+
830
+ # Offer to build the database here, since spawned workers cannot ask.
831
+ _offer_database_build(dbase_root)
832
+
833
+ # ---- dispatch: parallel for 2+ ions, serial otherwise ----
834
+ n_ions = len(worker_args)
835
+ if n_workers <= 0:
836
+ # Use the CPUs this process may run on, as its SLURM step allows, not
837
+ # the whole node's, which would start a worker per CPU of the node.
838
+ n_workers = (len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity")
839
+ else os.cpu_count() or 1)
840
+ use_parallel = n_workers > 1 and n_ions > 1
841
+
842
+ if use_parallel:
843
+ import multiprocessing as mp
844
+ pool_size = min(n_workers, n_ions)
845
+ ctx = mp.get_context("spawn")
846
+ print(f" Parallel: {pool_size} workers for {n_ions} ions")
847
+ with ctx.Pool(pool_size) as pool:
848
+ all_results = pool.map(_compute_single_ion, worker_args)
849
+ else:
850
+ all_results = [_compute_single_ion(a) for a in tqdm(
851
+ worker_args, desc="Computing G(T,N)", unit="ion"
852
+ )]
853
+
854
+ # ---- collect results ----
855
+ goft_dict: Dict[str, dict] = {}
856
+ for result in all_results:
857
+ for line_name, info in result.items():
858
+ used = info["hdf5_dbase_root"]
859
+ if dbase_root is not None and used != dbase_root:
860
+ raise RuntimeError(
861
+ f"A G(T,N) worker built its Ion against the CHIANTI "
862
+ f"database at {used} instead of the requested "
863
+ f"{dbase_root}, so the request did not reach it. Its "
864
+ f"contribution functions would come from the wrong "
865
+ f"atomic data."
866
+ )
867
+ print(
868
+ f" {line_name}: requested {info['target_wl_cm']*1e8:.4f} Angstrom, "
869
+ f"matched {info['matched_wl_aa']:.4f} Angstrom, "
870
+ f"{'observed' if info['observed'] else 'theoretical'} "
871
+ f"(delta={info['delta_aa']:.4f} Angstrom)"
872
+ )
873
+ # No line of the ion matches the name's wavelength to its digits.
874
+ if info["match"] == "nearest observed":
875
+ warnings.warn(
876
+ f"{line_name}: no line of this ion is at that wavelength to the digits "
877
+ f"given, so the nearest line CHIANTI has observed is used, at "
878
+ f"{info['matched_wl_aa']:.4f} Angstrom, {info['delta_aa']:.4f} Angstrom "
879
+ f"from the name, and synthesised there. To synthesise another line, "
880
+ f"theoretical ones included, give CHIANTI's wavelength for it.",
881
+ UserWarning, stacklevel=2)
882
+ elif info["match"] == "nearest theoretical":
883
+ warnings.warn(
884
+ f"{line_name}: no line of this ion is at that wavelength to the digits "
885
+ f"given, and CHIANTI has observed none of its lines, so the nearest of its "
886
+ f"theoretical wavelengths is used, at {info['matched_wl_aa']:.4f} Angstrom, "
887
+ f"{info['delta_aa']:.4f} Angstrom from the name, and synthesised there.",
888
+ UserWarning, stacklevel=2)
889
+ same = [other for other, seen in goft_dict.items()
890
+ if (seen["atom"], seen["ion"], seen["transition"])
891
+ == (info["atom"], info["ion"], info["transition"])]
892
+ if same:
893
+ raise ValueError(
894
+ f"{same[0]} and {line_name} are the same line, CHIANTI's at "
895
+ f"{info['matched_wl_aa']:.4f} Angstrom, which would be synthesised twice "
896
+ f"and summed. Give it once.")
897
+ goft_dict[line_name] = {
898
+ # Place the line at its CHIANTI wavelength, not at the
899
+ # name's digits.
900
+ "wl0": info["matched_wl_aa"] * u.AA.to(u.cm) * u.cm,
901
+ "transition": info["transition"],
902
+ "g_tn": info["g_tn"].astype(precision),
903
+ "atom": info["atom"],
904
+ "ion": info["ion"],
905
+ "hdf5_dbase_root": used,
906
+ }
907
+
908
+ return goft_dict, logT_grid.astype(precision), logN_grid.astype(precision)
909
+
910
+
911
+ ##############################################################################
912
+ # ---------------------------------------------------------------------------
913
+ # DEM and G(T) helpers
914
+ # ---------------------------------------------------------------------------
915
+ ##############################################################################
916
+
917
+ def compute_dem(
918
+ logT_cube: np.ndarray,
919
+ logN_cube: np.ndarray,
920
+ voxel_dh_cm: float | np.ndarray,
921
+ logT_grid: np.ndarray,
922
+ integration_axis: str = "z",
923
+ ) -> Tuple[np.ndarray, np.ndarray]:
924
+ """
925
+ DEM(T) and emission-measure-weighted mean electron density <n_e>(T).
926
+
927
+ Parameters
928
+ ----------
929
+ logT_cube : np.ndarray
930
+ 3D array of log10(T/K) values.
931
+ logN_cube : np.ndarray
932
+ 3D array of log10(n_e / cm^-3) values.
933
+ voxel_dh_cm : float or np.ndarray
934
+ Cell depth along the integration axis in cm: one value for all cells,
935
+ or one per cell along that axis.
936
+ logT_grid : np.ndarray
937
+ 1D array of temperature bin centers for DEM calculation.
938
+ integration_axis : str
939
+ View: the axis to integrate along, `"x"`, `"y"` or `"z"`, optionally
940
+ signed for the side of the box to look from, such as `"-x"`; see
941
+ `line_of_sight_velocity`.
942
+
943
+ Returns
944
+ -------
945
+ dem_map : np.ndarray
946
+ DEM array [cm^-5 per dex], with the two remaining spatial axes in
947
+ image order (row, column). Shape depends on the axis:
948
+ - "x": (nz, ny, nT)
949
+ - "y": (nz, nx, nT)
950
+ - "z": (ny, nx, nT)
951
+ avg_ne : np.ndarray
952
+ Mean electron density per T-bin [cm^-3]. Same shape as `dem_map`.
953
+ """
954
+ nT = len(logT_grid)
955
+ axis, _ = view_axis_and_side(integration_axis)
956
+
957
+ dlogT, logT_edges = _temperature_bins(logT_grid)
958
+
959
+ ne = 10.0 ** logN_cube.astype(np.float64)
960
+ dh = along_line_of_sight(voxel_dh_cm, axis)
961
+ w2 = np.broadcast_to(ne**2 * dh, logT_cube.shape) # weights for EM
962
+ w3 = np.broadcast_to(ne**3 * dh, logT_cube.shape) # weights for EM*n_e
963
+
964
+ # Split each cell between the two temperature bins either side of it, as
965
+ # in build_em_tv, so the DEM equals the emission measure build_em_tv
966
+ # spreads over velocity, and each bin's density is that of its plasma.
967
+ lower, share, inside = _neighbouring_bins(logT_cube, logT_grid, logT_edges)
968
+ pixel, (n_rows, n_cols) = _image_pixels(logT_cube.shape, integration_axis)
969
+ em = np.zeros(n_rows * n_cols * nT) # cm^-5
970
+ em_n = np.zeros_like(em) # cm^-5 * n_e
971
+ for step, of_T in ((0, 1.0 - share), (1, share)):
972
+ in_bin = (pixel * nT + np.minimum(lower + step, nT - 1)).ravel()
973
+ of_T = np.where(inside, of_T, 0.0)
974
+ np.add.at(em, in_bin, (w2 * of_T).ravel())
975
+ np.add.at(em_n, in_bin, (w3 * of_T).ravel())
976
+ em, em_n = em.reshape(n_rows, n_cols, nT), em_n.reshape(n_rows, n_cols, nT)
977
+
978
+ dem = em / dlogT
979
+ # Zero where there is no plasma at a temperature. Otherwise the division
980
+ # leaves uninitialised memory there, which ends up in the saved G
981
+ # diagnostic.
982
+ avg_ne = np.divide(em_n, em, out=np.zeros_like(em), where=em > 0.0)
983
+ return dem, avg_ne
984
+
985
+
986
+ def _temperature_bins(logT_grid: np.ndarray) -> Tuple[float, np.ndarray]:
987
+ """Width and edges of the DEM temperature bins centred on `logT_grid`."""
988
+ dlogT = logT_grid[1] - logT_grid[0] if len(logT_grid) > 1 else 0.1
989
+ logT_edges = np.concatenate([
990
+ [logT_grid[0] - dlogT/2],
991
+ logT_grid[:-1] + dlogT/2,
992
+ [logT_grid[-1] + dlogT/2]
993
+ ])
994
+ return dlogT, logT_edges
995
+
996
+
997
+ def _neighbouring_bins(values: np.ndarray, centres: np.ndarray,
998
+ edges: np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
999
+ """
1000
+ Split each value between the two bins either side of it.
1001
+
1002
+ Returns the index of the bin whose centre is at or below each value, the
1003
+ share going to the next bin up, and whether the value is inside the
1004
+ outer bin edges.
1005
+
1006
+ The share rises linearly from 0 at one centre to 1 at the next, so the
1007
+ two centres weighted by their shares average to the value: 2.4 km/s on a
1008
+ 5 km/s grid is 0.52 in the 0 km/s bin and 0.48 in the 5 km/s bin. A
1009
+ value between the first or last centre and that bin's outer edge goes
1010
+ wholly to that bin, so the grid still reaches its outer edges. A value
1011
+ beyond the edges is in no bin, and its share is zero.
1012
+ """
1013
+ values = np.asarray(values, dtype=np.float64)
1014
+ inside = (values >= edges[0]) & (values < edges[-1])
1015
+ n = centres.size
1016
+ lower = np.clip(np.searchsorted(centres, values, side="right") - 1, 0, max(n - 2, 0))
1017
+ if n > 1:
1018
+ share = np.clip((values - centres[lower]) / (centres[lower + 1] - centres[lower]),
1019
+ 0.0, 1.0)
1020
+ else:
1021
+ share = np.zeros(values.shape)
1022
+ return lower, np.where(inside, share, 0.0), inside
1023
+
1024
+
1025
+ def _image_pixels(shape: Tuple[int, int, int], integration_axis: str) -> Tuple[np.ndarray, Tuple[int, int]]:
1026
+ """
1027
+ Image pixel that each cell of a (z, y, x) cube is seen in.
1028
+
1029
+ For a view along `integration_axis`, returns each cell's flat pixel index
1030
+ (row times the number of columns plus column) and the image shape
1031
+ (rows, columns), in the order `compute_dem` and `build_em_tv` use.
1032
+
1033
+ Seen from the side opposite :data:`OBSERVER_SIDE`, the image is mirrored
1034
+ left to right, as walking round to the other side of the box mirrors it,
1035
+ so it is the right way round for that observer too.
1036
+ """
1037
+ axis, side = view_axis_and_side(integration_axis)
1038
+ z, y, x = np.indices(shape, sparse=True)
1039
+ rows, columns = {"x": (z, y), "y": (z, x), "z": (y, x)}[axis]
1040
+ n_rows, n_columns = rows.size, columns.size
1041
+ if side != OBSERVER_SIDE[axis]:
1042
+ columns = n_columns - 1 - columns
1043
+ return np.broadcast_to(rows * n_columns + columns, shape), (n_rows, n_columns)
1044
+
1045
+
1046
+ def _log_cubes(temperature: np.ndarray, electron_density: np.ndarray,
1047
+ precision: type) -> Tuple[np.ndarray, np.ndarray]:
1048
+ """log10 of the temperature and electron density, in `precision`.
1049
+
1050
+ A cell with no electrons gets a log density of minus infinity, so it
1051
+ adds nothing to the emission measure. A cell with no temperature gets
1052
+ log T = 0, below every temperature grid.
1053
+ """
1054
+ logN_cube = np.log10(electron_density, where=electron_density > 0.0,
1055
+ out=np.full_like(electron_density, -np.inf)).astype(precision)
1056
+ logT_cube = np.log10(temperature, where=temperature > 0.0,
1057
+ out=np.zeros_like(temperature)).astype(precision)
1058
+ return logT_cube, logN_cube
1059
+
1060
+
1061
+ def density_grid(
1062
+ temperature: np.ndarray,
1063
+ electron_density: np.ndarray,
1064
+ logT_grid: np.ndarray,
1065
+ precision: type = np.float64,
1066
+ ) -> Tuple[float, float, int]:
1067
+ """Density grid on which to compute G(T, n_e) for an atmosphere.
1068
+
1069
+ The contribution functions are read at the mean density of each
1070
+ temperature along each line of sight, and are zero off their grid. This
1071
+ grid uses the default grid's points (every 0.3 in log10 n_e, through
1072
+ 10**7 cm^-3) and extends each way to just past the densities of all
1073
+ cells whose temperature is on `logT_grid`. No mean density falls off
1074
+ it, and an atmosphere with a narrower density range than the default
1075
+ grid needs fewer points.
1076
+
1077
+ Parameters
1078
+ ----------
1079
+ temperature, electron_density : np.ndarray
1080
+ Temperature of every cell in K and electron density in cm^-3.
1081
+ logT_grid : np.ndarray
1082
+ DEM temperature grid, as returned by :func:`compute_goft_fiasco`.
1083
+ precision : type
1084
+ Precision of the synthesis; the densities are rounded to it as in
1085
+ :func:`synthesise_cubes`.
1086
+
1087
+ Returns
1088
+ -------
1089
+ logN_min, logN_max : float
1090
+ nN : int
1091
+ Grid for :func:`compute_goft_fiasco`. The default grid if no cell
1092
+ has electrons at a temperature on `logT_grid`.
1093
+ """
1094
+ points = _density_points(temperature, electron_density, logT_grid, precision)
1095
+ if points is None:
1096
+ return _LOGN_MIN, _LOGN_MAX, _N_N
1097
+ first, last = points
1098
+ return _density_point(first), _density_point(last), last - first + 1
1099
+
1100
+
1101
+ def _density_point(index: int) -> float:
1102
+ """log10 n_e at 0-based point `index` of the default density grid."""
1103
+ return _LOGN_MIN + index * (_LOGN_MAX - _LOGN_MIN) / (_N_N - 1)
1104
+
1105
+
1106
+ def _density_points(temperature: np.ndarray, electron_density: np.ndarray,
1107
+ logT_grid: np.ndarray, precision: type) -> Optional[Tuple[int, int]]:
1108
+ """First and last :func:`_density_point` indices of :func:`density_grid`.
1109
+
1110
+ None if no cell has electrons at a temperature on `logT_grid`.
1111
+ """
1112
+ logT_cube, logN_cube = _log_cubes(temperature, electron_density, precision)
1113
+ _, logT_edges = _temperature_bins(logT_grid)
1114
+ counted = ((logT_cube >= logT_edges[0]) & (logT_cube < logT_edges[-1])
1115
+ & np.isfinite(logN_cube))
1116
+ if not counted.any():
1117
+ return None
1118
+
1119
+ step = (_LOGN_MAX - _LOGN_MIN) / (_N_N - 1)
1120
+ # Positions of the lowest and highest densities, in steps from the
1121
+ # default grid's first point. Rounded so that a density on a point up to
1122
+ # floating-point error counts as on it.
1123
+ lowest, highest = (
1124
+ round((float(extreme) - _LOGN_MIN) / step, 9) for extreme in (
1125
+ np.min(logN_cube, where=counted, initial=np.inf),
1126
+ np.max(logN_cube, where=counted, initial=-np.inf)))
1127
+ # Take the points just past them, so a mean density a rounding error
1128
+ # beyond the lowest or highest is still on the grid.
1129
+ return int(np.ceil(lowest)) - 1, int(np.floor(highest)) + 1
1130
+
1131
+
1132
+ def interpolate_g_on_dem(
1133
+ goft: Dict[str, dict],
1134
+ avg_ne: np.ndarray,
1135
+ logT_grid: np.ndarray,
1136
+ logN_grid: np.ndarray,
1137
+ logT_goft: np.ndarray,
1138
+ precision: type = np.float32,
1139
+ ) -> None:
1140
+ """
1141
+ Each line's contribution function at each pixel's density and temperature.
1142
+
1143
+ G is interpolated linearly from the `compute_goft_fiasco` grid and is
1144
+ zero outside it. The result is stored in each line's entry of `goft` as
1145
+ `g`, shape (rows, columns, temperatures), for `synthesise_spectra`.
1146
+
1147
+ Parameters
1148
+ ----------
1149
+ goft : dict
1150
+ Lines from `compute_goft_fiasco`, modified in place.
1151
+ avg_ne : np.ndarray
1152
+ Electron density in each pixel at each temperature, in cm-3, shape
1153
+ (rows, columns, temperatures). If the cells along the line of sight
1154
+ differ in density, give their mean weighted by emission measure, as
1155
+ ECLIPSE's own synthesis does. For a DEM, give the density you choose.
1156
+ logT_grid : np.ndarray
1157
+ Temperatures at which to take G, as log10(T / K).
1158
+ logN_grid : np.ndarray
1159
+ Densities G was computed on, from `compute_goft_fiasco`.
1160
+ logT_goft : np.ndarray
1161
+ Temperatures G was computed on, from `compute_goft_fiasco`.
1162
+ precision : type, optional
1163
+ `np.float32` (default) or `np.float64`.
1164
+ """
1165
+ nT, n_rows, n_cols = len(logT_grid), *avg_ne.shape[:2]
1166
+
1167
+ # Build query points for interpolation
1168
+ logNe_flat = np.log10(avg_ne, where=avg_ne > 0.0,
1169
+ out=np.zeros_like(avg_ne)).transpose(2, 0, 1).ravel()
1170
+ logT_flat = np.broadcast_to(logT_grid[:, None, None],
1171
+ (nT, n_rows, n_cols)).ravel()
1172
+ query_pts = np.column_stack((logNe_flat, logT_flat))
1173
+
1174
+ for name, info in tqdm(goft.items(), desc="interpolating G", unit="line", leave=False):
1175
+ rgi = RegularGridInterpolator(
1176
+ (logN_grid, logT_goft), info["g_tn"],
1177
+ method="linear", bounds_error=False, fill_value=0.0
1178
+ )
1179
+ g_flat = rgi(query_pts)
1180
+ info["g"] = g_flat.reshape(nT, n_rows, n_cols).transpose(1, 2, 0).astype(precision)
1181
+
1182
+
1183
+ ##############################################################################
1184
+ # ---------------------------------------------------------------------------
1185
+ # Build EM(T,v) and synthesise spectra
1186
+ # ---------------------------------------------------------------------------
1187
+ ##############################################################################
1188
+
1189
+ # A view looks along one axis of the box from one side. A view named by its
1190
+ # axis alone looks from the side given by OBSERVER_SIDE (in utils): +1 for
1191
+ # the side of increasing coordinate, -1 for the other. From that side the
1192
+ # line cube, with rows and columns as create_line_cube lays them out, is the
1193
+ # right way round: the column axis crossed with the row axis points at the
1194
+ # observer, above the box (+z) for the top-down view and at +x and -y for
1195
+ # the two side views. A view named with a sign, such as "-x", looks from
1196
+ # that side. Seen from the side opposite OBSERVER_SIDE, the image is
1197
+ # mirrored left to right (_image_pixels), so it is the right way round for
1198
+ # that observer too.
1199
+ SIGNED_VIEWS = ("+x", "-x", "+y", "-y", "+z", "-z")
1200
+
1201
+
1202
+ def line_of_sight_velocity(velocity, integration_axis: str):
1203
+ """
1204
+ Velocity away from the observer, from velocity along the integration axis.
1205
+
1206
+ Simulation velocities are positive towards increasing coordinate, so an
1207
+ upflow has a positive z velocity. Seen from above, that upflow moves
1208
+ towards the observer and is blueshifted, so its line-of-sight velocity is
1209
+ negative. Positive line-of-sight velocity is a redshift: each line is
1210
+ placed at `lambda_0 (1 + v / c)`.
1211
+
1212
+ Parameters
1213
+ ----------
1214
+ velocity : np.ndarray or u.Quantity
1215
+ Velocity along the integration axis, positive towards increasing
1216
+ coordinate.
1217
+ integration_axis : str
1218
+ View: `"x"`, `"y"` or `"z"`, seen from the side given by
1219
+ :data:`OBSERVER_SIDE`, or a signed axis for the side, such as `"-x"`
1220
+ for the view from -x.
1221
+
1222
+ Returns
1223
+ -------
1224
+ np.ndarray or u.Quantity
1225
+ Velocity away from the observer, in the same units.
1226
+ """
1227
+ _, side = view_axis_and_side(integration_axis)
1228
+ return -side * velocity
1229
+
1230
+
1231
+ def build_em_tv(
1232
+ logT_cube: np.ndarray,
1233
+ vel_cube: np.ndarray,
1234
+ logT_grid: np.ndarray,
1235
+ vel_grid: np.ndarray,
1236
+ ne_sq_dh: np.ndarray,
1237
+ integration_axis: str = "z",
1238
+ ) -> np.ndarray:
1239
+ """
1240
+ 4-D emission measure cube EM(row, column, T, v) [cm^-5].
1241
+
1242
+ Each cell's emission measure is split between the two temperature bins
1243
+ and the two velocity bins either side of it, in proportion to how close
1244
+ it is to each (:func:`_neighbouring_bins`). Its mean velocity over the
1245
+ bins is then the cell's own, and its contribution function is
1246
+ interpolated in temperature. A value past the outermost bin centre but
1247
+ inside that bin goes wholly to it.
1248
+
1249
+ Parameters
1250
+ ----------
1251
+ logT_cube : np.ndarray
1252
+ 3D temperature cube, shape (nz, ny, nx).
1253
+ vel_cube : np.ndarray
1254
+ 3D velocity cube along the integration axis.
1255
+ logT_grid : np.ndarray
1256
+ Temperature bin centers.
1257
+ vel_grid : np.ndarray
1258
+ Velocity bin centers.
1259
+ ne_sq_dh : np.ndarray
1260
+ n_e^2 * dh for each voxel.
1261
+ integration_axis : str
1262
+ View: the axis to integrate along, `"x"`, `"y"` or `"z"`, optionally
1263
+ signed for the side of the box to look from, such as `"-x"`; see
1264
+ `line_of_sight_velocity`.
1265
+
1266
+ Returns
1267
+ -------
1268
+ em_tv : np.ndarray
1269
+ 4D emission measure cube, with the two remaining spatial axes in
1270
+ image order (row, column). Shape depends on the axis:
1271
+ - "x": (nz, ny, nT, nv)
1272
+ - "y": (nz, nx, nT, nv)
1273
+ - "z": (ny, nx, nT, nv)
1274
+ """
1275
+ axis, _ = view_axis_and_side(integration_axis)
1276
+ print(f" Building 4-D emission-measure cube along {axis}-axis...")
1277
+
1278
+ _, logT_edges = _temperature_bins(logT_grid)
1279
+ v_centres = (vel_grid.to_value(u.cm / u.s) if isinstance(vel_grid, u.Quantity)
1280
+ else np.asarray(vel_grid, dtype=float))
1281
+ v_edges = velocity_centers_to_edges(v_centres)
1282
+
1283
+ # Split each cell between the two temperature and two velocity bins
1284
+ # either side of it, in proportion to how close it is to each. Put wholly
1285
+ # in the nearest bin, a uniform 2.4 km/s flow on the default 5 km/s grid
1286
+ # would be synthesised at rest and a 2.6 km/s one at 5 km/s, and each
1287
+ # cell would get the contribution function of the nearest temperature
1288
+ # bin. With the split, the mean velocity over the bins is the cell's, the
1289
+ # line width grows by the spread of the two bins' velocities (at most a
1290
+ # quarter of a bin squared in variance), and the contribution function
1291
+ # is interpolated between the two temperatures. Past the outermost
1292
+ # centres, within the outermost bins, a cell goes to that bin alone.
1293
+ lower_T, share_T, in_T = _neighbouring_bins(logT_cube, logT_grid, logT_edges)
1294
+ lower_v, share_v, in_v = _neighbouring_bins(vel_cube, v_centres, v_edges)
1295
+
1296
+ # Plasma faster than the velocity grid emits outside the synthesised
1297
+ # wavelengths, so it is missing from the spectra but kept in the DEM.
1298
+ # Warn, since the default grid misses the fastest flows of a flare.
1299
+ in_temperature = ne_sq_dh * in_T
1300
+ beyond = in_temperature[~in_v].sum()
1301
+ if beyond > 0:
1302
+ warnings.warn(
1303
+ f"{100 * beyond / in_temperature.sum():.3g} per cent of the emission measure is "
1304
+ f"from plasma faster than the velocity grid reaches, "
1305
+ f"{v_edges[0] / 1e5:.1f} to {v_edges[-1] / 1e5:.1f} km/s, and emits beyond the "
1306
+ f"synthesised wavelengths, so it is not in the spectra. Widen --vel-lim to keep "
1307
+ f"it.", UserWarning, stacklevel=2)
1308
+
1309
+ # Sum along the line of sight into each image pixel's bins.
1310
+ nT, nv = len(logT_grid), len(v_centres)
1311
+ pixel, (n_rows, n_cols) = _image_pixels(logT_cube.shape, integration_axis)
1312
+ em = np.where(in_T & in_v, ne_sq_dh, 0.0)
1313
+ em_tv = np.zeros(n_rows * n_cols * nT * nv)
1314
+ for step_T, of_T in ((0, 1.0 - share_T), (1, share_T)):
1315
+ in_bin_T = pixel * nT + np.minimum(lower_T + step_T, nT - 1)
1316
+ for step_v, of_v in ((0, 1.0 - share_v), (1, share_v)):
1317
+ in_bin = in_bin_T * nv + np.minimum(lower_v + step_v, nv - 1)
1318
+ np.add.at(em_tv, in_bin.ravel(), (em * of_T * of_v).ravel())
1319
+ return em_tv.reshape(n_rows, n_cols, nT, nv)
1320
+
1321
+
1322
+ def synthesise_spectra(
1323
+ goft: Dict[str, dict],
1324
+ em_tv: np.ndarray,
1325
+ vel_grid: u.Quantity | np.ndarray,
1326
+ logT_grid: np.ndarray,
1327
+ ) -> None:
1328
+ """
1329
+ Spectrum of each line in every pixel, from the emission measure EM(T, v).
1330
+
1331
+ Each temperature and velocity bin adds a Gaussian with that
1332
+ temperature's thermal width, Doppler shifted by that velocity, scaled by
1333
+ the emission measure times G. The spectra are stored in each line's
1334
+ entry of `goft` as `si`, in erg / (s cm2 sr cm), on the wavelengths
1335
+ `wl_grid`: the velocity grid converted with `lambda_0 (1 + v / c)`.
1336
+
1337
+ Parameters
1338
+ ----------
1339
+ goft : dict
1340
+ Lines with `g` from `interpolate_g_on_dem`, modified in place.
1341
+ em_tv : np.ndarray
1342
+ Emission measure in each bin, in cm-5, shape (rows, columns,
1343
+ temperatures, velocities).
1344
+ vel_grid : u.Quantity or np.ndarray
1345
+ Velocity bin centres, evenly spaced and increasing. A plain array is
1346
+ taken to be in cm/s.
1347
+ logT_grid : np.ndarray
1348
+ Temperature bin centres, as log10(T / K).
1349
+ """
1350
+ kb = const.k_B.cgs.value
1351
+ c_cm_s = const.c.cgs.value
1352
+ # The Doppler shifts below are in cm/s. A grid in km/s taken as cm/s
1353
+ # would shift every line by a factor of 1e5 too little.
1354
+ vel_grid = (vel_grid.to(u.cm / u.s) if isinstance(vel_grid, u.Quantity)
1355
+ else np.asarray(vel_grid, dtype=float) * (u.cm / u.s))
1356
+
1357
+ # The wavelength grid below is the velocity grid mapped through
1358
+ # lambda_0 (1 + v/c), and create_line_cube sets CDELT from the first step
1359
+ # alone, so an uneven velocity grid would give a wrong wavelength axis.
1360
+ require_uniform_grid(vel_grid, "vel_grid")
1361
+
1362
+ for line, data in tqdm(goft.items(), desc="spectra", unit="line", leave=False):
1363
+ wl0 = data["wl0"].cgs.value # cm
1364
+
1365
+ # Create wavelength grid for this line
1366
+ data["wl_grid"] = (vel_grid * data["wl0"] / const.c + data["wl0"]).cgs
1367
+ wl_grid = data["wl_grid"].cgs.value # (n_lambda,)
1368
+
1369
+ _, atomic_weight = element_data(int(data["atom"]))
1370
+ atom_weight_g = (atomic_weight * u.u).cgs.value
1371
+
1372
+ # Thermal width per T-bin: sigma_T (nT,)
1373
+ sigma_T = wl0 * np.sqrt(kb * (10 ** logT_grid) / atom_weight_g) / c_cm_s
1374
+
1375
+ # Doppler-shifted center for each v-bin: (nv,)
1376
+ lam_cent = wl0 * (1 + vel_grid.value / c_cm_s)
1377
+
1378
+ # Build phi(T,v,lambda) as (nT,nv,n_lambda)
1379
+ delta = wl_grid[None, None, :] - lam_cent[None, :, None]
1380
+ phi = np.exp(-0.5 * (delta / sigma_T[:, None, None]) ** 2)
1381
+ phi /= sigma_T[:, None, None] * np.sqrt(2 * np.pi)
1382
+
1383
+ # EM(row,col,T,v) * G(T) -> (n_rows,n_cols,nT,nv)
1384
+ weighted = em_tv * data["g"][..., None]
1385
+
1386
+ # Collapse T and v: dot ((nT,nv) , (nT,nv)) -> (n_rows,n_cols,n_lambda)
1387
+ spec_map = np.tensordot(weighted, phi, axes=([2, 3], [0, 1]))
1388
+
1389
+ data["si"] = spec_map / (4 * np.pi)
1390
+
1391
+
1392
+ def synthesise_cubes(
1393
+ temperature: np.ndarray,
1394
+ electron_density: np.ndarray,
1395
+ los_velocity: np.ndarray,
1396
+ dh_cm,
1397
+ goft: Dict[str, dict],
1398
+ logT_grid: np.ndarray,
1399
+ logN_grid: np.ndarray,
1400
+ vel_grid: u.Quantity,
1401
+ integration_axis: str,
1402
+ precision: type,
1403
+ *,
1404
+ continuum: Optional[Dict[str, dict]] = None,
1405
+ return_density: bool = False,
1406
+ ) -> Union[Tuple[Dict[str, dict], np.ndarray, np.ndarray],
1407
+ Tuple[Dict[str, dict], np.ndarray, np.ndarray, np.ndarray]]:
1408
+ """
1409
+ Spectra of every line from one set of cubes (a whole box or a strip of it).
1410
+
1411
+ Parameters
1412
+ ----------
1413
+ temperature : np.ndarray
1414
+ Temperature of every cell in K, shape `(nz, ny, nx)`.
1415
+ electron_density : np.ndarray
1416
+ Electron density of every cell in cm^-3, same shape.
1417
+ los_velocity : np.ndarray
1418
+ Velocity away from the observer in cm/s, same shape.
1419
+ dh_cm : float or np.ndarray
1420
+ Cell depth along the line of sight in cm, one value for all cells or
1421
+ one per cell along that axis.
1422
+ goft : dict
1423
+ Contribution functions from :func:`compute_goft_fiasco`, on the
1424
+ temperature grid `logT_grid` and density grid `logN_grid`. Not
1425
+ modified: the entries are copied before the interpolated G and the
1426
+ spectra are added.
1427
+ logT_grid, logN_grid : np.ndarray
1428
+ Grids the contribution functions are tabulated on. The DEM uses the
1429
+ same temperature grid.
1430
+ vel_grid : u.Quantity
1431
+ Velocity bin centres, evenly spaced.
1432
+ integration_axis : str
1433
+ View: `"x"`, `"y"` or `"z"`, optionally signed for the side of the
1434
+ box to look from, such as `"-x"`; see `line_of_sight_velocity`.
1435
+ precision : type
1436
+ np.float32 or np.float64.
1437
+ continuum : dict, optional
1438
+ Continuum to add, by entry name, each with `"wl_grid"` (its
1439
+ wavelengths) and `"free"` and `"two_photon"` from
1440
+ `continuum.compute_continuum_fiasco` on these temperatures and
1441
+ densities. Keyword only. Default None (no continuum).
1442
+ return_density : bool, optional
1443
+ Also return the electron density at which the contribution functions
1444
+ are taken: for each temperature of each pixel, the mean over its
1445
+ cells weighted by emission measure. Keyword only. Default False.
1446
+
1447
+ Returns
1448
+ -------
1449
+ lines : dict
1450
+ Copy of `goft` whose entries also hold `"g"` (contribution function
1451
+ on the DEM), `"wl_grid"` and `"si"` (specific intensity, shape
1452
+ `(rows, columns, wavelength)`). With `continuum`, it also has an
1453
+ entry per continuum name, with `"wl_grid"`, `"si"` and `"wl0"` (the
1454
+ middle of its wavelengths), and no `"atom"` or `"ion"`.
1455
+ dem_map : np.ndarray
1456
+ As returned by :func:`compute_dem`.
1457
+ em_tv : np.ndarray
1458
+ As returned by :func:`build_em_tv`.
1459
+ electron_density : np.ndarray
1460
+ Only with `return_density`: density in cm^-3, same shape as
1461
+ `dem_map`, zero where there is no plasma at a temperature.
1462
+ """
1463
+ logT_cube, logN_cube = _log_cubes(temperature, electron_density, precision)
1464
+
1465
+ dem_map, avg_ne_map = compute_dem(logT_cube, logN_cube, dh_cm, logT_grid, integration_axis)
1466
+
1467
+ # The contribution functions are zero off their density grid; warn, with
1468
+ # the fraction of the emission measure affected.
1469
+ occupied = dem_map > 0
1470
+ log_ne = np.log10(avg_ne_map, where=avg_ne_map > 0,
1471
+ out=np.full(avg_ne_map.shape, -np.inf))
1472
+ off_grid = occupied & ((log_ne < logN_grid[0]) | (log_ne > logN_grid[-1]))
1473
+ if off_grid.any():
1474
+ warnings.warn(
1475
+ f"{100 * dem_map[off_grid].sum() / dem_map[occupied].sum():.3g} per cent of the "
1476
+ f"emission measure is at electron densities outside the {10 ** logN_grid[0]:.0e} "
1477
+ f"to {10 ** logN_grid[-1]:.0e} cm^-3 the contribution functions are worked out "
1478
+ f"for, where they are taken as zero.", UserWarning, stacklevel=2)
1479
+
1480
+ lines = {name: dict(info) for name, info in goft.items()}
1481
+ interpolate_g_on_dem(lines, avg_ne_map, logT_grid, logN_grid, logT_grid, precision)
1482
+
1483
+ ne_sq_dh = ((10.0 ** logN_cube.astype(np.float64)) ** 2
1484
+ * along_line_of_sight(dh_cm, view_axis_and_side(integration_axis)[0]))
1485
+ em_tv = build_em_tv(logT_cube, los_velocity, logT_grid, vel_grid, ne_sq_dh, integration_axis)
1486
+
1487
+ synthesise_spectra(lines, em_tv, vel_grid, logT_grid)
1488
+
1489
+ # The continuum uses the emission measure in each temperature bin from
1490
+ # the DEM, and each pixel's density at that temperature. The DEM keeps
1491
+ # cells faster than the velocity grid, which the continuum needs since it
1492
+ # is not Doppler shifted.
1493
+ if continuum:
1494
+ emission_measure = dem_map * _temperature_bins(logT_grid)[0]
1495
+ for name, table in continuum.items():
1496
+ grid = u.Quantity(table["wl_grid"]).to(u.cm)
1497
+ spectra = continuum_spectra(emission_measure, avg_ne_map, logN_grid,
1498
+ table["free"], table["two_photon"])
1499
+ lines[name] = {"si": spectra.astype(precision), "wl_grid": grid,
1500
+ "wl0": grid[grid.size // 2], "atom": None, "ion": None}
1501
+ if return_density:
1502
+ return lines, dem_map, em_tv, avg_ne_map
1503
+ return lines, dem_map, em_tv
1504
+
1505
+
1506
+ def _cell_size_mm(cube: NDCube, pixel_axis: int) -> float:
1507
+ """World size in Mm of one pixel of `cube` along `pixel_axis`.
1508
+
1509
+ Measured between the edges of the first pixel, so it needs neither a
1510
+ second pixel nor a CDELT, which a cropped cube's WCS does not have.
1511
+ """
1512
+ # A cropped cube wraps its WCS in a high-level object; the pixel-to-world
1513
+ # conversion of plain numbers is on the low-level one underneath.
1514
+ wcs = getattr(cube.wcs, "low_level_wcs", cube.wcs)
1515
+ low = [0.0] * wcs.pixel_n_dim
1516
+ high = [0.0] * wcs.pixel_n_dim
1517
+ low[pixel_axis], high[pixel_axis] = -0.5, 0.5
1518
+ world_low = wcs.pixel_to_world_values(*low)
1519
+ world_high = wcs.pixel_to_world_values(*high)
1520
+ unit = u.Unit(wcs.world_axis_units[pixel_axis])
1521
+ return ((world_high[pixel_axis] - world_low[pixel_axis]) * unit).to_value(u.Mm)
1522
+
1523
+
1524
+ def _world_at(coords: u.Quantity, crpix: float) -> float:
1525
+ """
1526
+ Value of the uniform grid `coords` at 1-based pixel `crpix`, as a float.
1527
+
1528
+ A reference pixel in the middle of an axis with an even number of pixels
1529
+ falls between two of them, so its value is read off the grid there
1530
+ rather than taken from the pixel below.
1531
+ """
1532
+ if coords.size == 1:
1533
+ return coords[0].value
1534
+ step = (coords[-1] - coords[0]).value / (coords.size - 1)
1535
+ return coords[0].value + (crpix - 1) * step
1536
+
1537
+
1538
+ def create_line_cube(
1539
+ line_name: str,
1540
+ line_data: dict,
1541
+ spatial_cube: NDCube,
1542
+ intensity_unit: u.Unit,
1543
+ integration_axis: str = "z",
1544
+ ) -> NDCube:
1545
+ """
1546
+ NDCube of one line's synthesised spectra, with the image's coordinates.
1547
+
1548
+ Parameters
1549
+ ----------
1550
+ line_name : str
1551
+ Line name.
1552
+ line_data : dict
1553
+ The line's entry from `compute_goft_fiasco`, with its spectra `si`
1554
+ and their wavelengths `wl_grid` added by `synthesise_spectra`.
1555
+ spatial_cube : NDCube
1556
+ Atmosphere cube, such as from `create_atmosphere_ndcube`, whose
1557
+ coordinates the image takes.
1558
+ intensity_unit : u.Unit
1559
+ Unit of the spectra.
1560
+ integration_axis : str, optional
1561
+ View used in the synthesis: the axis looked along, `"x"`, `"y"` or
1562
+ `"z"` (default), optionally signed for the side of the box looked
1563
+ from, such as `"-x"`; see `line_of_sight_velocity`. Seen from the
1564
+ side opposite :data:`OBSERVER_SIDE`, the image is mirrored left to
1565
+ right, so its column coordinate is the box coordinate with its sign
1566
+ flipped, which increases to that observer's right.
1567
+
1568
+ Returns
1569
+ -------
1570
+ NDCube
1571
+ Spectra indexed `[row, column, wavelength]` like an image, so summing
1572
+ over wavelength gives an array that plots the right way up.
1573
+ """
1574
+ view = view_name(integration_axis)
1575
+ los_axis, side = view_axis_and_side(view)
1576
+ # An axis with unequal cells has no single CDELT. Only the line of sight
1577
+ # may be like that, since it is integrated out here.
1578
+ nonuniform = (spatial_cube.meta or {}).get("nonuniform_axes", [])
1579
+ stretched = [axis for axis in AXES
1580
+ if axis != los_axis and axis in nonuniform]
1581
+ if stretched:
1582
+ raise ValueError(
1583
+ f"The {', '.join(stretched)} axis of the atmosphere is not evenly "
1584
+ f"spaced, so it cannot be an image axis of a view along "
1585
+ f"{los_axis}. Only the line of sight may be stretched.")
1586
+
1587
+ # The simulation cubes are (z, y, x), so integrating one axis out leaves
1588
+ # 'si' already in (row, column, wavelength) order for every view.
1589
+ cube_data = line_data["si"]
1590
+
1591
+ # Take each axis's cell size from the reference cube's WCS, so an axis
1592
+ # one cell wide has one too. It is read as the world size of one pixel,
1593
+ # which works for any WCS, cropped ones included.
1594
+ cell_size = [_cell_size_mm(spatial_cube, pixel_axis) for pixel_axis in range(3)]
1595
+
1596
+ # The WCS below has a single linear CDELT, so the wavelength grid must be
1597
+ # uniform. Checked here as well as in synthesise_spectra because this is
1598
+ # a public entry point, called directly by the DEM and VDEM routes.
1599
+ wl_step = require_uniform_grid(line_data["wl_grid"].to(u.cm), "wl_grid")
1600
+
1601
+ # Get spatial coordinate information from the reference cube,
1602
+ # whose array axes are (z, y, x)
1603
+ if los_axis == "x":
1604
+ # Integration along X -> data shape (nz, ny, n_lambda): rows are Z, columns are Y
1605
+ nz, ny, nl = cube_data.shape
1606
+ y_coords = spatial_cube.axis_world_coords(1)[0] # Y coordinates
1607
+ z_coords = spatial_cube.axis_world_coords(0)[0] # Z coordinates
1608
+
1609
+ spatial_axes = ['WAVE', 'SOLY', 'SOLZ'] # Wavelength, Y, Z
1610
+ spatial_units = ['cm', 'Mm', 'Mm']
1611
+ spatial_cdelt = [
1612
+ wl_step,
1613
+ cell_size[1],
1614
+ cell_size[2],
1615
+ ]
1616
+ spatial_crpix = [(nl + 1) / 2, (ny + 1) / 2, 1] # Wavelength centered, Y centered, Z at first pixel
1617
+ spatial_crval = [
1618
+ _world_at(line_data["wl_grid"].to(u.cm), spatial_crpix[0]),
1619
+ _world_at(y_coords.to(u.Mm), spatial_crpix[1]),
1620
+ z_coords[0].to(u.Mm).value # Z starts where original cube starts
1621
+ ]
1622
+
1623
+ elif los_axis == "y":
1624
+ # Integration along Y -> data shape (nz, nx, n_lambda): rows are Z, columns are X
1625
+ nz, nx, nl = cube_data.shape
1626
+ x_coords = spatial_cube.axis_world_coords(2)[0] # X coordinates
1627
+ z_coords = spatial_cube.axis_world_coords(0)[0] # Z coordinates
1628
+
1629
+ spatial_axes = ['WAVE', 'SOLX', 'SOLZ'] # Wavelength, X, Z
1630
+ spatial_units = ['cm', 'Mm', 'Mm']
1631
+ spatial_cdelt = [
1632
+ wl_step,
1633
+ cell_size[0],
1634
+ cell_size[2],
1635
+ ]
1636
+ spatial_crpix = [(nl + 1) / 2, (nx + 1) / 2, 1] # Wavelength centered, X centered, Z at first pixel
1637
+ spatial_crval = [
1638
+ _world_at(line_data["wl_grid"].to(u.cm), spatial_crpix[0]),
1639
+ _world_at(x_coords.to(u.Mm), spatial_crpix[1]),
1640
+ z_coords[0].to(u.Mm).value # Z starts where original cube starts
1641
+ ]
1642
+
1643
+ else: # los_axis == "z"
1644
+ # Integration along Z -> data shape (ny, nx, n_lambda): rows are Y, columns are X
1645
+ ny, nx, nl = cube_data.shape
1646
+ x_coords = spatial_cube.axis_world_coords(2)[0] # X coordinates
1647
+ y_coords = spatial_cube.axis_world_coords(1)[0] # Y coordinates
1648
+
1649
+ spatial_axes = ['WAVE', 'SOLX', 'SOLY'] # Wavelength, X, Y
1650
+ spatial_units = ['cm', 'Mm', 'Mm']
1651
+ spatial_cdelt = [
1652
+ wl_step,
1653
+ cell_size[0],
1654
+ cell_size[1],
1655
+ ]
1656
+ spatial_crpix = [(nl + 1) / 2, (nx + 1) / 2, (ny + 1) / 2] # All centered
1657
+ spatial_crval = [
1658
+ _world_at(line_data["wl_grid"].to(u.cm), spatial_crpix[0]),
1659
+ _world_at(x_coords.to(u.Mm), spatial_crpix[1]),
1660
+ _world_at(y_coords.to(u.Mm), spatial_crpix[2]),
1661
+ ]
1662
+
1663
+ # Seen from the other side, the image is mirrored left to right
1664
+ # (_image_pixels), so the column coordinate is the box coordinate with
1665
+ # its sign flipped. The columns are centred, so the middle one stays put.
1666
+ if side != OBSERVER_SIDE[los_axis]:
1667
+ spatial_crval[1] = -spatial_crval[1]
1668
+
1669
+ wcs = WCS(naxis=3)
1670
+ wcs.wcs.ctype = spatial_axes
1671
+ wcs.wcs.cunit = spatial_units
1672
+ wcs.wcs.crpix = spatial_crpix
1673
+ wcs.wcs.crval = spatial_crval
1674
+ wcs.wcs.cdelt = spatial_cdelt
1675
+
1676
+ return NDCube(
1677
+ cube_data,
1678
+ wcs=wcs,
1679
+ unit=intensity_unit,
1680
+ meta={
1681
+ "line_name": line_name,
1682
+ "rest_wav": line_data["wl0"],
1683
+ "atom": line_data["atom"],
1684
+ "ion": line_data["ion"],
1685
+ "integration_axis": view,
1686
+ "velocity_convention": VELOCITY_CONVENTION,
1687
+ "spatial_reference": spatial_cube.meta if hasattr(spatial_cube, 'meta') else None
1688
+ }
1689
+ )
1690
+
1691
+
1692
+
1693
+ ##############################################################################
1694
+ # ---------------------------------------------------------------------------
1695
+ # M A I N W O R K F L O W
1696
+ # ---------------------------------------------------------------------------
1697
+ ##############################################################################
1698
+
1699
+ # Options giving the location and layout of MURaM's own files, used by
1700
+ # dynamic mode and by the deprecated static route without an atmosphere
1701
+ # file. An atmosphere file carries all of this itself, so giving both is an
1702
+ # error.
1703
+ MURAM_LAYOUT_OPTIONS = ("data_dir", "temp_file", "rho_file", "vx_file", "vy_file",
1704
+ "vz_file", "cube_shape", "voxel_dx", "voxel_dy", "voxel_dz")
1705
+
1706
+ # Options read only in dynamic mode, which a static synthesis from an
1707
+ # atmosphere file would ignore.
1708
+ DYNAMIC_OPTIONS = ("slit_width", "temp_dir", "temp_filename", "rho_dir", "rho_filename",
1709
+ "vx_dir", "vx_filename", "vy_dir", "vy_filename", "vz_dir",
1710
+ "vz_filename", "time_dir", "time_filename")
1711
+
1712
+ # Documentation of the atmosphere file and how to write one.
1713
+ ATMOSPHERE_DOCS = "https://solarc-eclipse.readthedocs.io/en/stable/synthesis/#atmosphere-files"
1714
+ # Documentation of observing a time series of atmosphere files, which
1715
+ # replaces the deprecated dynamic mode.
1716
+ TIME_SERIES_DOCS = "https://solarc-eclipse.readthedocs.io/en/stable/time-series/"
1717
+
1718
+
1719
+ class _NotedOption(argparse.Action):
1720
+ """Store the value and record that the option was given, even at its default."""
1721
+
1722
+ def __call__(self, parser, namespace, values, option_string=None):
1723
+ setattr(namespace, self.dest, values)
1724
+ given = getattr(namespace, "given_options", None)
1725
+ if given is None:
1726
+ given = set()
1727
+ namespace.given_options = given
1728
+ given.add(self.dest)
1729
+
1730
+
1731
+ def build_parser() -> argparse.ArgumentParser:
1732
+ """Command line parser for synthesise-spectra."""
1733
+ parser = argparse.ArgumentParser(
1734
+ description="Synthesise solar spectra from 3D MHD simulation data",
1735
+ formatter_class=argparse.ArgumentDefaultsHelpFormatter
1736
+ )
1737
+
1738
+ # Input/Output paths
1739
+ parser.add_argument("--atmosphere", type=str, default=None,
1740
+ help="The ECLIPSE atmosphere file (HDF5) to synthesise "
1741
+ f"from; see {ATMOSPHERE_DOCS} for how to write one. "
1742
+ "Required, except in dynamic mode.")
1743
+ # MURaM file options of the deprecated static route, hidden from the
1744
+ # help so old command lines still run, with a warning, until the route
1745
+ # is removed.
1746
+ for flag, default in (("--temp-file", "temp/eosT.0270000"),
1747
+ ("--rho-file", "rho/result_prim_0.0270000"),
1748
+ ("--vx-file", "vx/result_prim_1.0270000"),
1749
+ ("--vy-file", "vy/result_prim_3.0270000"),
1750
+ ("--vz-file", "vz/result_prim_2.0270000")):
1751
+ parser.add_argument(flag, type=str, default=default, action=_NotedOption,
1752
+ help=argparse.SUPPRESS)
1753
+ parser.add_argument("--output-dir", type=str, default="./run/input",
1754
+ help="Output directory for results")
1755
+ parser.add_argument("--output-name", type=str, default="synthesised_spectra.h5",
1756
+ help="Output filename, an HDF5 synthesis file; a name ending in .pkl "
1757
+ "writes the deprecated pickle instead")
1758
+
1759
+ # Line / abundance specification (fiasco)
1760
+ parser.add_argument("--lines", nargs="+", required=True,
1761
+ help="Line specifications (e.g. Fe12_195.1190 Fe09_171.073)")
1762
+ parser.add_argument("--abundance", type=str, default="sun_coronal_2021_chianti",
1763
+ help="CHIANTI abundance dataset name for fiasco")
1764
+ parser.add_argument("--continuum", action="store_true",
1765
+ help="Also synthesise the free-free, free-bound and two-photon "
1766
+ "continuum under the lines' windows, as entries of its own "
1767
+ "beside the lines")
1768
+ parser.add_argument("--n-workers", type=int, default=0,
1769
+ help="Number of parallel workers for fiasco G(T,N) "
1770
+ "computation (0 = all CPUs, default: 0)")
1771
+ parser.add_argument("--hdf5-dbase-root", type=str, default=None,
1772
+ help="CHIANTI HDF5 database to use, overriding the one "
1773
+ "in ~/.fiasco/fiascorc. Use this to run against a "
1774
+ "second CHIANTI version without changing the "
1775
+ "default for other work.")
1776
+ parser.add_argument("--goft-temperature-chunk", type=int, default=None,
1777
+ help="Compute G(T,N) this many temperatures at a time "
1778
+ "rather than the whole grid at once. This lowers "
1779
+ "fiasco's peak memory, which for an ion with many "
1780
+ "levels can be several GB, at the cost of reading "
1781
+ "the atomic data again for each chunk. The result "
1782
+ "is the same.")
1783
+
1784
+ # Integration direction
1785
+ parser.add_argument("--integration-axis", choices=["x", "y", "z", *SIGNED_VIEWS],
1786
+ default="z",
1787
+ help="The axis to look along, x, y or z, from +x, -y or above; or the "
1788
+ "axis with a sign for the side to look from, such as -x")
1789
+
1790
+ # Cropping parameters (in Heliocentric coordinates)
1791
+ parser.add_argument("--crop-x", nargs=2, type=str, default=None,
1792
+ help="Crop in x direction: x_min x_max (e.g. '-50 Mm' '50 Mm')")
1793
+ parser.add_argument("--crop-y", nargs=2, type=str, default=None,
1794
+ help="Crop in y direction: y_min y_max (e.g. '-50 Mm' '50 Mm')")
1795
+ parser.add_argument("--crop-z", nargs=2, type=str, default=None,
1796
+ help="Crop in z direction: z_min z_max (e.g. '0 Mm' '20 Mm')")
1797
+
1798
+ # Velocity grid
1799
+ parser.add_argument("--vel-res", type=str, default="5.0 km/s",
1800
+ help="Velocity resolution (e.g. '5.0 km/s')")
1801
+ parser.add_argument("--vel-lim", type=str, default="300.0 km/s",
1802
+ help="Velocity limit +/- (e.g. '300.0 km/s')")
1803
+
1804
+ # Processing options
1805
+ parser.add_argument("--downsample", type=int, default=1,
1806
+ help="Downsampling factor (1 = no downsampling)")
1807
+ parser.add_argument("--precision", choices=["float32", "float64"], default="float64",
1808
+ help="Numerical precision")
1809
+ parser.add_argument("--mass-per-electron", "--mean-mol-wt",
1810
+ dest="mass_per_electron", type=float, default=None,
1811
+ help="Mass of the plasma per free electron, in atomic mass "
1812
+ "units, which turns a mass density into an electron "
1813
+ "density. By default it is worked out from --abundance "
1814
+ "for a fully ionised plasma, about 1.16 for coronal "
1815
+ "abundances. --mean-mol-wt is the old name; ECLIPSE "
1816
+ "0.11.0 and earlier used 1.29, the value for a neutral gas. "
1817
+ "Not used when the atmosphere gives an electron density.")
1818
+
1819
+ # Dynamic atmosphere mode (time-varying synthesis). It reads MURaM's own
1820
+ # files, so it carries the options describing their layout.
1821
+ dynamic_group = parser.add_argument_group("Dynamic atmosphere mode (deprecated)",
1822
+ "Options for synthesising with time-varying atmosphere (raster scanning). "
1823
+ "Deprecated: observe a time series of atmosphere files with the "
1824
+ "instrument run instead.")
1825
+ dynamic_group.add_argument("--slit-rest-time", type=str, default=None,
1826
+ help="Slit rest time per position (e.g. '40 s'). "
1827
+ "Enables dynamic mode when specified.")
1828
+ dynamic_group.add_argument("--slit-width", type=str, default=None,
1829
+ action=_NotedOption,
1830
+ help="Slit width (e.g. '0.2 arcsec', required for dynamic mode)")
1831
+ dynamic_group.add_argument("--data-dir", type=str, default="data/atmosphere",
1832
+ action=_NotedOption,
1833
+ help="Directory containing the MURaM files")
1834
+ dynamic_group.add_argument("--cube-shape", nargs=3, type=int, default=[512, 768, 256],
1835
+ action=_NotedOption,
1836
+ help="Cube dimensions in the file's storage order (nx nz ny)")
1837
+ dynamic_group.add_argument("--voxel-dx", type=str, default="0.192 Mm",
1838
+ action=_NotedOption,
1839
+ help="Voxel size in x (e.g. '0.192 Mm')")
1840
+ dynamic_group.add_argument("--voxel-dy", type=str, default="0.192 Mm",
1841
+ action=_NotedOption,
1842
+ help="Voxel size in y (e.g. '0.192 Mm')")
1843
+ dynamic_group.add_argument("--voxel-dz", type=str, default="0.064 Mm",
1844
+ action=_NotedOption,
1845
+ help="Voxel size in z (e.g. '0.064 Mm')")
1846
+
1847
+ # Directory arguments for dynamic mode
1848
+ dynamic_group.add_argument("--temp-dir", type=str, default=None,
1849
+ action=_NotedOption,
1850
+ help="Directory containing temperature files (for dynamic mode)")
1851
+ dynamic_group.add_argument("--temp-filename", type=str, default="eosT",
1852
+ action=_NotedOption,
1853
+ help="Temperature filename prefix before timestep suffix")
1854
+ dynamic_group.add_argument("--rho-dir", type=str, default=None,
1855
+ action=_NotedOption,
1856
+ help="Directory containing density files (for dynamic mode)")
1857
+ dynamic_group.add_argument("--rho-filename", type=str, default="result_prim_0",
1858
+ action=_NotedOption,
1859
+ help="Density filename prefix before timestep suffix")
1860
+ dynamic_group.add_argument("--vx-dir", type=str, default=None,
1861
+ action=_NotedOption,
1862
+ help="Directory containing vx files (for dynamic mode)")
1863
+ dynamic_group.add_argument("--vx-filename", type=str, default="result_prim_1",
1864
+ action=_NotedOption,
1865
+ help="Vx filename prefix before timestep suffix")
1866
+ dynamic_group.add_argument("--vy-dir", type=str, default=None,
1867
+ action=_NotedOption,
1868
+ help="Directory containing vy files (for dynamic mode)")
1869
+ dynamic_group.add_argument("--vy-filename", type=str, default="result_prim_3",
1870
+ action=_NotedOption,
1871
+ help="Vy filename prefix before timestep suffix")
1872
+ dynamic_group.add_argument("--vz-dir", type=str, default=None,
1873
+ action=_NotedOption,
1874
+ help="Directory containing vz files (for dynamic mode)")
1875
+ dynamic_group.add_argument("--vz-filename", type=str, default="result_prim_2",
1876
+ action=_NotedOption,
1877
+ help="Vz filename prefix before timestep suffix")
1878
+ dynamic_group.add_argument("--time-dir", type=str, default="header",
1879
+ action=_NotedOption,
1880
+ help="Directory containing header files (for dynamic mode)")
1881
+ dynamic_group.add_argument("--time-filename", type=str, default="Header",
1882
+ action=_NotedOption,
1883
+ help="Header filename prefix before timestep suffix")
1884
+
1885
+ return parser
1886
+
1887
+
1888
+ def parse_arguments(argv=None):
1889
+ """Parse command line arguments for spectrum synthesis."""
1890
+ words = list(sys.argv[1:] if argv is None else argv)
1891
+ # argparse takes a value starting with a dash, such as the -x in
1892
+ # "--integration-axis -x", for an option and stops. Joined into
1893
+ # "--integration-axis=-x", it is read as the value.
1894
+ joined, i = [], 0
1895
+ while i < len(words):
1896
+ if (words[i] == "--integration-axis" and i + 1 < len(words)
1897
+ and words[i + 1] in SIGNED_VIEWS):
1898
+ joined.append(f"--integration-axis={words[i + 1]}")
1899
+ i += 2
1900
+ else:
1901
+ joined.append(words[i])
1902
+ i += 1
1903
+ return build_parser().parse_args(joined)
1904
+
1905
+
1906
+ def check_atmosphere_options(args) -> None:
1907
+ """
1908
+ Reject options that --atmosphere makes meaningless; warn if it is absent.
1909
+
1910
+ The synthesis reads its atmosphere from an atmosphere file. Without one,
1911
+ static mode reads MURaM's own files and dynamic mode reads a time series
1912
+ of them; both are deprecated. The MURaM layout options describe those
1913
+ files and the dynamic mode options apply only to dynamic mode, so any of
1914
+ them given with --atmosphere would be silently ignored.
1915
+ """
1916
+ if not args.atmosphere:
1917
+ if args.slit_rest_time is not None:
1918
+ warnings.warn(
1919
+ f"Dynamic mode (--slit-rest-time) is deprecated and will be "
1920
+ f"removed in a future release: write the snapshots as atmosphere "
1921
+ f"files and observe them as a time series in the instrument run, "
1922
+ f"as described at {TIME_SERIES_DOCS}, with 'direction: decreasing' "
1923
+ f"under 'raster:' to scan the way dynamic mode does.",
1924
+ FutureWarning, stacklevel=2)
1925
+ else:
1926
+ warnings.warn(
1927
+ f"No --atmosphere was given, so the synthesis is reading "
1928
+ f"MURaM's own files from {args.data_dir}. This is deprecated "
1929
+ f"and will be removed in a future release: write the snapshot "
1930
+ f"as an atmosphere file, as described at {ATMOSPHERE_DOCS}, "
1931
+ f"and give it with --atmosphere.",
1932
+ FutureWarning, stacklevel=2)
1933
+ return
1934
+ # The parser notes every layout option that appeared on the command
1935
+ # line, so one typed at its default value is caught too.
1936
+ noted = getattr(args, "given_options", ())
1937
+ given = [name for name in MURAM_LAYOUT_OPTIONS if name in noted]
1938
+ if given:
1939
+ flags = ", ".join("--" + name.replace("_", "-") for name in given)
1940
+ raise ValueError(
1941
+ f"--atmosphere carries the cube shape, cell sizes and data itself, "
1942
+ f"so {flags} would not be used. Those options describe MURaM's "
1943
+ f"own files.")
1944
+ if args.slit_rest_time is not None:
1945
+ raise ValueError(
1946
+ f"Dynamic mode reads its time series from MURaM files and cannot "
1947
+ f"take an atmosphere file. A time series of atmosphere files is "
1948
+ f"observed by the instrument run instead: see {TIME_SERIES_DOCS}.")
1949
+ given = [name for name in DYNAMIC_OPTIONS if name in noted]
1950
+ if given:
1951
+ flags = ", ".join("--" + name.replace("_", "-") for name in given)
1952
+ raise ValueError(
1953
+ f"--atmosphere is a static synthesis, so {flags} would not be "
1954
+ f"used. Those options only apply to dynamic mode (--slit-rest-time).")
1955
+
1956
+
1957
+ def load_atmosphere_file(
1958
+ path: str | Path,
1959
+ integration_axis: str,
1960
+ downsample: int | bool = False,
1961
+ crop_x=None, crop_y=None, crop_z=None,
1962
+ ) -> Atmosphere:
1963
+ """
1964
+ Read an atmosphere file, with the velocity along `integration_axis`.
1965
+
1966
+ Downsampling and cropping are applied here, in the atmosphere's own
1967
+ coordinates. The two image axes are checked to be evenly spaced, since
1968
+ the image WCS can only describe an even grid. The line of sight may be
1969
+ stretched: the emission measure uses the true size of each cell along it.
1970
+ """
1971
+ atmosphere = read_atmosphere(path, velocities=(integration_axis,))
1972
+ return _prepare_atmosphere(atmosphere, str(path), integration_axis,
1973
+ downsample, crop_x, crop_y, crop_z)
1974
+
1975
+
1976
+ def load_muram_files(
1977
+ args,
1978
+ integration_axis: str,
1979
+ downsample: int | bool = False,
1980
+ crop_x=None, crop_y=None, crop_z=None,
1981
+ ) -> Atmosphere:
1982
+ """
1983
+ Read the MURaM files named by the deprecated static options.
1984
+
1985
+ They are downsampled, cropped and checked as in
1986
+ :func:`load_atmosphere_file`. x and y are centred on zero and z = 0 is at
1987
+ the centre of the bottom cell, as ECLIPSE has always placed a MURaM box,
1988
+ so the crop options keep the meaning they had before atmosphere files.
1989
+ """
1990
+ data_dir = Path(args.data_dir)
1991
+ files = {
1992
+ "temperature": (args.temp_file, u.K),
1993
+ "mass_density": (args.rho_file, u.g / u.cm**3),
1994
+ f"velocity_{integration_axis}": (getattr(args, f"v{integration_axis}_file"),
1995
+ u.cm / u.s),
1996
+ }
1997
+ cubes = {}
1998
+ for name, (file_name, unit) in files.items():
1999
+ path = data_dir / file_name
2000
+ if not path.exists():
2001
+ raise FileNotFoundError(f"{name} file not found: {path}")
2002
+ cubes[name] = load_cube(path, shape=tuple(args.cube_shape), unit=unit)
2003
+
2004
+ nz, ny, nx = cubes["temperature"].shape
2005
+ voxel = {axis: u.Quantity(getattr(args, f"voxel_d{axis}")) for axis in AXES}
2006
+ atmosphere = Atmosphere(
2007
+ x_edges=(np.arange(nx + 1) - nx / 2) * voxel["x"],
2008
+ y_edges=(np.arange(ny + 1) - ny / 2) * voxel["y"],
2009
+ z_edges=(np.arange(nz + 1) - 0.5) * voxel["z"],
2010
+ source="MURaM", **cubes)
2011
+ return _prepare_atmosphere(atmosphere, f"the MURaM files in {data_dir}",
2012
+ integration_axis, downsample, crop_x, crop_y, crop_z)
2013
+
2014
+
2015
+ def _prepare_atmosphere(atmosphere: Atmosphere, name: str, integration_axis: str,
2016
+ downsample, crop_x, crop_y, crop_z) -> Atmosphere:
2017
+ """Downsample and crop `atmosphere`, and check its image axes are evenly spaced."""
2018
+ if downsample:
2019
+ atmosphere = atmosphere.downsampled(downsample)
2020
+ if crop_x or crop_y or crop_z:
2021
+ atmosphere = atmosphere.cropped(x=crop_x, y=crop_y, z=crop_z)
2022
+ for axis in AXES:
2023
+ if axis != integration_axis and not atmosphere.is_uniform(axis):
2024
+ raise ValueError(
2025
+ f"The {axis} axis of {name} is not evenly spaced, and with the "
2026
+ f"line of sight along {integration_axis} it would become an "
2027
+ f"image axis, whose coordinates must be even. Only the axis "
2028
+ f"along the line of sight may be stretched; resample the "
2029
+ f"others onto an even grid first.")
2030
+ return atmosphere
2031
+
2032
+
2033
+ def resolve_mass_per_electron(args) -> Tuple[float, str]:
2034
+ """Mass per free electron in atomic mass units, and where it came from."""
2035
+ given = getattr(args, "mass_per_electron", None)
2036
+ # Accept the old name, mean_mol_wt, from scripts that set it on `args`.
2037
+ if given is None and getattr(args, "mean_mol_wt", None) is not None:
2038
+ warnings.warn(
2039
+ "mean_mol_wt is the old name of mass_per_electron, the mass of the plasma per "
2040
+ "free electron in atomic mass units; it is used, but set mass_per_electron "
2041
+ "instead.", FutureWarning, stacklevel=2)
2042
+ given = args.mean_mol_wt
2043
+ if given is not None:
2044
+ try:
2045
+ value = require_mass_per_electron(given)
2046
+ except ValueError as error:
2047
+ raise ValueError(f"--mass-per-electron: {error}") from None
2048
+ return value, "given on the command line"
2049
+ value = mass_per_electron(args.abundance, getattr(args, "hdf5_dbase_root", None))
2050
+ return value, f"fully ionised plasma with {args.abundance} abundances"
2051
+
2052
+
2053
+ def along_line_of_sight(values, integration_axis: str) -> np.ndarray:
2054
+ """
2055
+ Reshape `values` to broadcast over a (z, y, x) cube.
2056
+
2057
+ `values` holds one value per cell along the line of sight. A single
2058
+ value is returned unchanged, so one cell size applies everywhere.
2059
+ """
2060
+ values = np.asarray(values, dtype=float)
2061
+ if values.ndim == 0:
2062
+ return values
2063
+ if values.ndim != 1:
2064
+ raise ValueError(f"Expected one value per cell along the line of sight, "
2065
+ f"got an array of shape {values.shape}.")
2066
+ shape = [1, 1, 1]
2067
+ shape[NUMPY_AXIS[integration_axis]] = values.size
2068
+ return values.reshape(shape)
2069
+
2070
+
2071
+ def main(args=None) -> None:
2072
+ """
2073
+ Synthesise solar spectra from a 3D MHD simulation.
2074
+
2075
+ Supports two modes:
2076
+ - Static mode: one timestep, from an atmosphere file (--atmosphere) or
2077
+ from MURaM's own files (deprecated).
2078
+ - Dynamic mode (deprecated): time-varying synthesis with raster
2079
+ scanning, from MURaM's own files. The instrument run observes a time
2080
+ series of atmosphere files instead (see eclipse.raster).
2081
+
2082
+ Parameters
2083
+ ----------
2084
+ args : argparse.Namespace, optional
2085
+ Command line arguments. If None, parsed from sys.argv.
2086
+ """
2087
+ if args is None:
2088
+ args = parse_arguments()
2089
+
2090
+ # ---------------- Configuration from arguments -----------------
2091
+ precision = np.float32 if args.precision == "float32" else np.float64
2092
+ if args.downsample < 1:
2093
+ raise ValueError(f"--downsample must be 1 or more, got {args.downsample}.")
2094
+ downsample = args.downsample if args.downsample > 1 else False
2095
+ # Check the line names now, before reading the atmosphere, which can take
2096
+ # minutes and gigabytes.
2097
+ for name in args.lines:
2098
+ _parse_line_name(name)
2099
+ vel_res = u.Quantity(args.vel_res)
2100
+ vel_lim = u.Quantity(args.vel_lim)
2101
+ # Check the velocity grid now, before the atmosphere is read or anything
2102
+ # computed.
2103
+ velocity_grid(vel_res, vel_lim, ("--vel-res", "--vel-lim"))
2104
+
2105
+ intensity_unit = u.erg/u.s/u.cm**2/u.sr/u.cm
2106
+
2107
+ print_mem = lambda: f"{psutil.virtual_memory().used/1e9:.2f}/" \
2108
+ f"{psutil.virtual_memory().total/1e9:.2f} GB"
2109
+
2110
+ # The view: the axis to read and sum the atmosphere along, and the side
2111
+ # of the box it is seen from.
2112
+ view = view_name(args.integration_axis.lower())
2113
+ integration_axis, observer_side = view_axis_and_side(view)
2114
+ check_atmosphere_options(args)
2115
+ # A name ending in .pkl still gets the old pickle format, so older scripts
2116
+ # keep working until it is removed.
2117
+ write_pickle = Path(args.output_name).suffix.lower() in (".pkl", ".pickle", ".dill")
2118
+ if write_pickle:
2119
+ warnings.warn(
2120
+ f"--output-name {args.output_name} names a pickle, as older versions of ECLIPSE "
2121
+ f"wrote the synthesis. Writing one is deprecated and will stop in a future "
2122
+ f"release: give the output a name ending in .h5 for a synthesis file.",
2123
+ FutureWarning, stacklevel=2)
2124
+ elif Path(args.output_name).suffix.lower() not in (".h5", ".hdf5", ".hdf"):
2125
+ warnings.warn(
2126
+ f"--output-name {args.output_name} gets a synthesis file, which is HDF5. Older "
2127
+ f"versions of ECLIPSE wrote a pickle under any name; a name ending in .pkl "
2128
+ f"still gets one, until that is removed.", UserWarning, stacklevel=2)
2129
+
2130
+ # Inputs to the common processing below, set by whichever route reads
2131
+ # the atmosphere. Only an atmosphere file can give the electron density
2132
+ # directly, or a different size for each cell along the line of sight;
2133
+ # dynamic mode gives a mass density and one cell size.
2134
+ ne_values = None
2135
+ los_thickness = None
2136
+ atmosphere_metadata = None
2137
+
2138
+ # Determine if we're in dynamic mode
2139
+ dynamic_mode = args.slit_rest_time is not None
2140
+
2141
+ if dynamic_mode:
2142
+ # Validate dynamic mode requirements
2143
+ if args.slit_width is None:
2144
+ raise ValueError("--slit-width is required for dynamic mode (when --slit-rest-time is specified)")
2145
+ # The snapshots are laid out along x, which the slit steps across in a
2146
+ # view along z or y. In a view along x, x is the line of sight, and
2147
+ # every pixel would sum cells from all the snapshots.
2148
+ if integration_axis == "x":
2149
+ raise ValueError(
2150
+ "Dynamic mode lays its snapshots across x, which is the line of sight "
2151
+ "of a view along x, so every pixel would add up all of them. View along "
2152
+ "z or y.")
2153
+ if view != integration_axis:
2154
+ raise ValueError(
2155
+ f"Dynamic mode, which is deprecated, looks only from the side of the box "
2156
+ f"that --integration-axis {integration_axis} names. To look from "
2157
+ f"{view}, synthesise from an atmosphere file with --atmosphere.")
2158
+
2159
+ base_dir = Path(args.data_dir)
2160
+ # Voxel sizes in the MURaM files. load_cube scales them itself when it
2161
+ # downsamples, so they are passed to it as given.
2162
+ file_voxel_dz = u.Quantity(args.voxel_dz)
2163
+ file_voxel_dx = u.Quantity(args.voxel_dx)
2164
+ file_voxel_dy = u.Quantity(args.voxel_dy)
2165
+
2166
+ # Voxel sizes of the cubes as synthesised, for the path length along
2167
+ # the line of sight, the slice timing and the saved metadata.
2168
+ voxel_dz = file_voxel_dz * (downsample or 1)
2169
+ voxel_dx = file_voxel_dx * (downsample or 1)
2170
+ voxel_dy = file_voxel_dy * (downsample or 1)
2171
+
2172
+ # Parse slit rest time and slit width
2173
+ slit_rest_time = u.Quantity(args.slit_rest_time)
2174
+ slit_width = u.Quantity(args.slit_width)
2175
+
2176
+ # Determine which velocity direction to use
2177
+ if integration_axis == "x":
2178
+ vel_dir = args.vx_dir or "vx"
2179
+ vel_filename = args.vx_filename
2180
+ voxel_dh = voxel_dx
2181
+ elif integration_axis == "y":
2182
+ vel_dir = args.vy_dir or "vy"
2183
+ vel_filename = args.vy_filename
2184
+ voxel_dh = voxel_dy
2185
+ else: # "z"
2186
+ vel_dir = args.vz_dir or "vz"
2187
+ vel_filename = args.vz_filename
2188
+ voxel_dh = voxel_dz
2189
+
2190
+ # Set directory defaults
2191
+ temp_dir = args.temp_dir or "temp"
2192
+ rho_dir = args.rho_dir or "rho"
2193
+
2194
+ print(f"DYNAMIC MODE - Time-varying synthesis at MHD resolution (deprecated)")
2195
+ print(f" Slit width: {slit_width}")
2196
+ print(f" Slit rest time: {slit_rest_time}")
2197
+ print(f" Voxel dx: {voxel_dx}")
2198
+ print()
2199
+
2200
+ # Discover available timesteps
2201
+ time_dir = base_dir / args.time_dir
2202
+ print(f"Discovering timesteps from {time_dir}...")
2203
+ timestep_times = discover_timesteps(time_dir, args.time_filename)
2204
+ print(f" Found {len(timestep_times)} timesteps")
2205
+ for suffix, sim_time in sorted(timestep_times.items(), key=lambda x: x[1]):
2206
+ print(f" {suffix}: {sim_time:.3f} s")
2207
+ print()
2208
+
2209
+ # Calculate MHD cube dimensions
2210
+ cube_shape_tuple = tuple(args.cube_shape)
2211
+ nx_mhd = cube_shape_tuple[0]
2212
+ if downsample:
2213
+ nx_mhd = nx_mhd // downsample
2214
+
2215
+ # Prepare crop_x for slice mapping if specified
2216
+ crop_x_for_mapping = None
2217
+ if args.crop_x:
2218
+ crop_x_for_mapping = (u.Quantity(args.crop_x[0]), u.Quantity(args.crop_x[1]))
2219
+
2220
+ # Compute slice-to-timestep mapping at MHD resolution
2221
+ print(f"Computing slice-to-timestep mapping at MHD resolution...")
2222
+ slice_mapping, grouped_slices = compute_slice_timestep_mapping_mhd(
2223
+ nx_mhd, voxel_dx, slit_width, slit_rest_time, timestep_times,
2224
+ crop_x=crop_x_for_mapping
2225
+ )
2226
+ print(f" MHD slices per timestep:")
2227
+ for suffix, indices in sorted(grouped_slices.items(), key=lambda x: min(x[1])):
2228
+ print(f" {suffix}: {len(indices)} slices (indices {min(indices)}-{max(indices)})")
2229
+ print()
2230
+
2231
+ # Build composite cubes at MHD resolution
2232
+ print(f"Building composite atmosphere cubes at MHD resolution ({print_mem()})...")
2233
+ temp_cube, rho_cube, vel_cube = build_composite_cubes_mhd(
2234
+ base_dir=base_dir,
2235
+ temp_dir=temp_dir,
2236
+ temp_filename=args.temp_filename,
2237
+ rho_dir=rho_dir,
2238
+ rho_filename=args.rho_filename,
2239
+ vel_dir=vel_dir,
2240
+ vel_filename=vel_filename,
2241
+ slice_mapping=slice_mapping,
2242
+ grouped_slices=grouped_slices,
2243
+ cube_shape=cube_shape_tuple,
2244
+ voxel_dx=file_voxel_dx,
2245
+ voxel_dy=file_voxel_dy,
2246
+ voxel_dz=file_voxel_dz,
2247
+ downsample=downsample,
2248
+ precision=precision,
2249
+ )
2250
+ print(f" Composite cube shape: {temp_cube.data.shape}")
2251
+
2252
+ # Apply cropping if requested in dynamic mode
2253
+ if args.crop_x or args.crop_y or args.crop_z:
2254
+ print(f"Applying cropping ({print_mem()})")
2255
+ temp_cube, rho_cube, vel_cube = apply_cube_cropping(
2256
+ temp_cube, rho_cube, vel_cube,
2257
+ args.crop_x, args.crop_y, args.crop_z
2258
+ )
2259
+ print(f" Cropped cubes to shape: {temp_cube.data.shape}")
2260
+
2261
+ reference_cube = temp_cube
2262
+ rho = u.Quantity(rho_cube.data, rho_cube.unit)
2263
+
2264
+ # Dynamic mode metadata for output (no spatial rebinning in synthesis)
2265
+ dynamic_mode_metadata = {
2266
+ "enabled": True,
2267
+ "slit_width": slit_width,
2268
+ "slit_rest_time": slit_rest_time,
2269
+ "scan_direction": "right_to_left",
2270
+ "slice_timesteps": slice_mapping,
2271
+ "available_timesteps": timestep_times,
2272
+ "spatially_rebinned": False, # Output is at MHD resolution
2273
+ }
2274
+
2275
+ else:
2276
+ # Static mode: from an atmosphere file, which carries its own layout,
2277
+ # or by the deprecated route from MURaM's own files, read into the
2278
+ # same atmosphere the converter would write.
2279
+ dynamic_mode_metadata = {"enabled": False}
2280
+
2281
+ if args.atmosphere:
2282
+ print("STATIC MODE - Synthesis from an atmosphere file")
2283
+ print(f" Atmosphere: {args.atmosphere}")
2284
+ else:
2285
+ print("STATIC MODE - Synthesis from MURaM files (deprecated)")
2286
+ print(f" Data directory: {args.data_dir}")
2287
+ print(f" Integration axis: {view}")
2288
+ print(f" Velocity grid: +/-{vel_lim:.1f} at {vel_res:.1f} resolution")
2289
+ print(f" Precision: {precision}")
2290
+ if downsample:
2291
+ print(f" Downsampling: {downsample}x")
2292
+ print(f" Lines: {args.lines}")
2293
+ print(f" Abundance: {args.abundance}")
2294
+ if args.crop_x or args.crop_y or args.crop_z:
2295
+ print(f" Cropping: X={args.crop_x}, Y={args.crop_y}, Z={args.crop_z}")
2296
+ print()
2297
+
2298
+ print(f"Reading the atmosphere ({print_mem()})")
2299
+ crops = dict(crop_x=args.crop_x, crop_y=args.crop_y, crop_z=args.crop_z)
2300
+ if args.atmosphere:
2301
+ atmosphere = load_atmosphere_file(
2302
+ args.atmosphere, integration_axis, downsample=downsample, **crops)
2303
+ else:
2304
+ atmosphere = load_muram_files(
2305
+ args, integration_axis, downsample=downsample, **crops)
2306
+ print(atmosphere.describe())
2307
+
2308
+ # The file may hold float32 in any units. Convert to the --precision
2309
+ # dtype and to K and cm/s, which the processing below assumes.
2310
+ temp_cube = atmosphere.to_ndcube(
2311
+ atmosphere.temperature.astype(precision).to(u.K))
2312
+ vel_cube = atmosphere.to_ndcube(
2313
+ atmosphere.velocity(integration_axis).astype(precision).to(u.cm / u.s))
2314
+ rho = None
2315
+ if atmosphere.mass_density is not None:
2316
+ rho = atmosphere.mass_density.astype(precision)
2317
+ if atmosphere.electron_density is not None:
2318
+ ne_values = atmosphere.electron_density.astype(precision).to_value(u.cm**-3)
2319
+ los_thickness = atmosphere.cell_thickness(integration_axis)
2320
+ reference_cube = temp_cube
2321
+
2322
+ # The cell sizes are saved in place of the MURaM voxel sizes; a
2323
+ # stretched axis has no single size.
2324
+ voxel_dx, voxel_dy, voxel_dz = (
2325
+ atmosphere.spacing(axis) if atmosphere.is_uniform(axis) else None
2326
+ for axis in AXES)
2327
+ atmosphere_metadata = {
2328
+ "path": str(Path(args.atmosphere).resolve()) if args.atmosphere else None,
2329
+ "source": atmosphere.source,
2330
+ "time": atmosphere.time,
2331
+ "shape": atmosphere.shape,
2332
+ "nonuniform_axes": atmosphere.nonuniform_axes(),
2333
+ "electron_density_given": atmosphere.electron_density is not None,
2334
+ }
2335
+
2336
+ # ---------------- Common processing (both modes) -----------------
2337
+
2338
+ # Build velocity grid
2339
+ vel_grid = velocity_grid(vel_res, vel_lim, ("--vel-res", "--vel-lim"))
2340
+
2341
+ # Electron density from the atmosphere if given, otherwise the mass
2342
+ # density divided by the mass per free electron.
2343
+ if ne_values is None:
2344
+ mass_per_electron_amu, mass_per_electron_source = resolve_mass_per_electron(args)
2345
+ print(f"Electron density from the mass density with "
2346
+ f"{mass_per_electron_amu:.4f} u per electron ({mass_per_electron_source})")
2347
+ ne_values = (rho / (mass_per_electron_amu * const.u)).to_value(u.cm**-3)
2348
+ else:
2349
+ mass_per_electron_amu = None
2350
+ mass_per_electron_source = "not needed: the atmosphere gives the electron density"
2351
+ print("Electron density taken from the atmosphere")
2352
+
2353
+ # The velocity files hold the velocity along each axis; the Doppler shift
2354
+ # needs the velocity away from the observer.
2355
+ vel_data = line_of_sight_velocity(vel_cube.data, view)
2356
+
2357
+ # ---------------- Compute contribution functions (fiasco) ---------
2358
+ # Only over the densities this atmosphere has.
2359
+ logN_min, logN_max, nN = density_grid(
2360
+ temp_cube.data, ne_values,
2361
+ np.linspace(_LOGT_MIN, _LOGT_MAX, _N_T).astype(precision), precision)
2362
+ print(f"Computing contribution functions via fiasco at 10^{logN_min:.1f} to "
2363
+ f"10^{logN_max:.1f} cm^-3 ({print_mem()})")
2364
+ goft, logT_goft, logN_grid = compute_goft_fiasco(
2365
+ args.lines, abundance=args.abundance, precision=precision,
2366
+ logN_min=logN_min, logN_max=logN_max, nN=nN,
2367
+ n_workers=args.n_workers,
2368
+ hdf5_dbase_root=getattr(args, "hdf5_dbase_root", None),
2369
+ temperature_chunk=getattr(args, "goft_temperature_chunk", None),
2370
+ )
2371
+
2372
+ # Record the database actually used, not the one requested, so a run that
2373
+ # did not choose one can still be traced to its atomic data.
2374
+ goft_dbase_root = (
2375
+ next(iter(goft.values()))["hdf5_dbase_root"] if goft else None
2376
+ )
2377
+ print(f" CHIANTI database: {goft_dbase_root}")
2378
+
2379
+ # Use the GOFT temperature grid as our DEM temperature grid
2380
+ logT_grid = logT_goft
2381
+
2382
+ # The size of each cell along the line of sight. MURaM's files have one
2383
+ # size; an atmosphere file may give every cell its own.
2384
+ if los_thickness is None:
2385
+ los_thickness = {"x": voxel_dx, "y": voxel_dy, "z": voxel_dz}[integration_axis]
2386
+ dh_cm = los_thickness.to_value(u.cm)
2387
+
2388
+ # ---------------- DEM, EM(T,v) and spectra -----------------
2389
+ # Continuum over the lines' windows, one entry per group of overlapping
2390
+ # windows. Each window is the line's velocity grid around its CHIANTI
2391
+ # wavelength, as in synthesise_spectra.
2392
+ continuum = None
2393
+ if getattr(args, "continuum", False):
2394
+ windows = continuum_windows([(vel_grid * info["wl0"] / const.c + info["wl0"]).cgs
2395
+ for info in goft.values()])
2396
+ print(f"Computing the continuum over {len(windows)} window"
2397
+ f"{'' if len(windows) == 1 else 's'} via fiasco ({print_mem()})")
2398
+ wavelength = np.concatenate([grid.to_value(u.cm) for grid in windows.values()]) * u.cm
2399
+ free, two_photon = compute_continuum_fiasco(
2400
+ wavelength, logT_grid, logN_grid, abundance=args.abundance,
2401
+ n_workers=args.n_workers, hdf5_dbase_root=goft_dbase_root)
2402
+ continuum, start = {}, 0
2403
+ for name, grid in windows.items():
2404
+ stop = start + grid.size
2405
+ continuum[name] = {"wl_grid": grid, "free": free[:, start:stop],
2406
+ "two_photon": two_photon[:, :, start:stop]}
2407
+ start = stop
2408
+
2409
+ print(f"Calculating the DEM, the emission measure in (T,v) and the spectra ({print_mem()})")
2410
+ goft, dem_map, em_tv, electron_density = synthesise_cubes(
2411
+ temp_cube.data, ne_values, vel_data, dh_cm, goft, logT_grid, logN_grid,
2412
+ vel_grid, view, precision, continuum=continuum, return_density=True)
2413
+
2414
+ # ---------------- Create output cubes -----------------
2415
+ print(f"Creating output cubes ({print_mem()})")
2416
+ line_cubes = {}
2417
+ for name, info in goft.items():
2418
+ line_cubes[name] = create_line_cube(
2419
+ name, info, reference_cube, intensity_unit, view
2420
+ )
2421
+
2422
+ print(f"Built {len(line_cubes)} line cubes")
2423
+
2424
+ # ---------------- Save results -----------------
2425
+ output_dir = Path(args.output_dir)
2426
+ output_dir.mkdir(parents=True, exist_ok=True)
2427
+ output_file = output_dir / args.output_name
2428
+
2429
+ # The line cubes are the file's spectra. The other synthesis products are
2430
+ # saved with them, so the file holds the whole synthesis.
2431
+ products = {
2432
+ "dem_map": dem_map,
2433
+ "em_tv": em_tv,
2434
+ # Electron density at each temperature of each pixel (mean over its
2435
+ # cells weighted by emission measure), at which G was taken.
2436
+ "electron_density": electron_density,
2437
+ "logT_grid": logT_grid,
2438
+ "vel_grid": vel_grid,
2439
+ "logN_grid": logN_grid,
2440
+ # Each line's spectra and wavelengths are already in the line cubes,
2441
+ # so they are not stored twice.
2442
+ "goft": {name: {key: value for key, value in info.items()
2443
+ if key not in ("si", "wl_grid")}
2444
+ for name, info in goft.items() if name not in (continuum or {})},
2445
+ "voxel_sizes": {"dx": voxel_dx, "dy": voxel_dy, "dz": voxel_dz},
2446
+ "dynamic_mode": dynamic_mode_metadata,
2447
+ "atmosphere": atmosphere_metadata,
2448
+ "config": {
2449
+ "precision": precision.__name__,
2450
+ "downsample": downsample,
2451
+ "vel_res": vel_res,
2452
+ "vel_lim": vel_lim,
2453
+ "mass_per_electron": mass_per_electron_amu,
2454
+ "mass_per_electron_source": mass_per_electron_source,
2455
+ "intensity_unit": str(intensity_unit),
2456
+ "atmosphere": args.atmosphere,
2457
+ "cube_shape": None if args.atmosphere else args.cube_shape,
2458
+ "data_dir": None if args.atmosphere else str(Path(args.data_dir)),
2459
+ "lines": args.lines,
2460
+ "abundance": args.abundance,
2461
+ "continuum": bool(getattr(args, "continuum", False)),
2462
+ "hdf5_dbase_root": goft_dbase_root,
2463
+ "integration_axis": view,
2464
+ "velocity_convention": VELOCITY_CONVENTION,
2465
+ "observer_side": observer_side,
2466
+ "crop_params": {
2467
+ "crop_x": args.crop_x,
2468
+ "crop_y": args.crop_y,
2469
+ "crop_z": args.crop_z
2470
+ }
2471
+ }
2472
+ }
2473
+
2474
+ if write_pickle:
2475
+ # The old pickle format, including the full contribution functions.
2476
+ with open(output_file, "wb") as f:
2477
+ dill.dump({"line_cubes": line_cubes, **products,
2478
+ "goft": {name: info for name, info in goft.items()
2479
+ if name not in (continuum or {})}}, f)
2480
+ else:
2481
+ # Save the snapshot time, so syntheses of a series of snapshots can be
2482
+ # observed as a time series.
2483
+ write_line_cubes(line_cubes, output_file,
2484
+ source=(atmosphere_metadata or {}).get("source") or "",
2485
+ products=products,
2486
+ time=(atmosphere_metadata or {}).get("time"))
2487
+
2488
+ print(f"Saved results to {output_file} ({os.path.getsize(output_file) / 1e6:.2f} MB)")
2489
+ print("Synthesis complete!")
2490
+
2491
+ if __name__ == "__main__":
2492
+ main()