| 1 | ### script to download opacities from DACE |
|---|
| 2 | ### adapted from https://github.com/bmorris3/shone/tree/refs/heads/opacity-downloads |
|---|
| 3 | |
|---|
| 4 | import logging |
|---|
| 5 | import warnings |
|---|
| 6 | from functools import cached_property |
|---|
| 7 | import os |
|---|
| 8 | import tarfile |
|---|
| 9 | import shutil |
|---|
| 10 | |
|---|
| 11 | import numpy as np |
|---|
| 12 | import xarray as xr |
|---|
| 13 | from astropy.table import Table |
|---|
| 14 | from dace_query.opacity import Molecule, Atom |
|---|
| 15 | |
|---|
| 16 | from chemistry import species_name_to_common_isotopologue_name |
|---|
| 17 | |
|---|
| 18 | interp_kwargs = dict( |
|---|
| 19 | method='nearest', |
|---|
| 20 | kwargs=dict(fill_value="extrapolate") |
|---|
| 21 | ) |
|---|
| 22 | |
|---|
| 23 | __all__ = [ |
|---|
| 24 | 'download_molecule', |
|---|
| 25 | 'download_atom' |
|---|
| 26 | ] |
|---|
| 27 | |
|---|
| 28 | class AvailableOpacities: |
|---|
| 29 | @cached_property |
|---|
| 30 | def atomic(self): |
|---|
| 31 | return get_atomic_database() |
|---|
| 32 | |
|---|
| 33 | @cached_property |
|---|
| 34 | def molecular(self): |
|---|
| 35 | return get_molecular_database() |
|---|
| 36 | |
|---|
| 37 | def get_atomic_database_entry(self, atom, charge, line_list): |
|---|
| 38 | table = self.atomic |
|---|
| 39 | return table[( |
|---|
| 40 | (table['atom'] == atom) & |
|---|
| 41 | (table['line_list'] == line_list) & |
|---|
| 42 | (table['charge'] == charge) |
|---|
| 43 | )] |
|---|
| 44 | |
|---|
| 45 | def get_molecular_database_entry(self, isotopologue, line_list): |
|---|
| 46 | table = self.molecular |
|---|
| 47 | return table[( |
|---|
| 48 | (table['isotopologue'] == isotopologue) & |
|---|
| 49 | (table['line_list'] == line_list) |
|---|
| 50 | )] |
|---|
| 51 | |
|---|
| 52 | def get_molecular_line_lists(self, isotopologue): |
|---|
| 53 | table = self.molecular |
|---|
| 54 | return set(table[table['isotopologue'] == isotopologue]['line_list']) |
|---|
| 55 | |
|---|
| 56 | def get_atomic_line_lists(self, atom): |
|---|
| 57 | table = self.atomic |
|---|
| 58 | return set(table[table['atom'] == atom]['line_list']) |
|---|
| 59 | |
|---|
| 60 | def get_atomic_latest_version(self, atom, charge, line_list): |
|---|
| 61 | table = self.atomic |
|---|
| 62 | matches = table[( |
|---|
| 63 | (table['atom'] == atom) & |
|---|
| 64 | (table['line_list'] == line_list) & |
|---|
| 65 | (table['charge'] == charge) |
|---|
| 66 | )] |
|---|
| 67 | return max(set(matches['version'])) |
|---|
| 68 | |
|---|
| 69 | def get_molecular_latest_version(self, isotopologue, line_list): |
|---|
| 70 | table = self.molecular |
|---|
| 71 | matches = table[( |
|---|
| 72 | (table['isotopologue'] == isotopologue) & |
|---|
| 73 | (table['line_list'] == line_list) |
|---|
| 74 | )] |
|---|
| 75 | return max(set(matches['version'])) |
|---|
| 76 | |
|---|
| 77 | def get_atomic_pT_range(self, atom, charge, line_list): |
|---|
| 78 | table = self.get_atomic_database_entry(atom, charge, line_list) |
|---|
| 79 | temperature_range = ( |
|---|
| 80 | int(table['temp_min_k'][0]), |
|---|
| 81 | int(table['temp_max_k'][0]) |
|---|
| 82 | ) |
|---|
| 83 | pressure_range = ( |
|---|
| 84 | float(table['pressure_min_exponent_b'][0]), |
|---|
| 85 | float(table['pressure_max_exponent_b'][0]) |
|---|
| 86 | ) |
|---|
| 87 | return temperature_range, pressure_range |
|---|
| 88 | |
|---|
| 89 | def get_molecular_pT_range(self, isotopologue, line_list): |
|---|
| 90 | table = self.get_molecular_database_entry(isotopologue, line_list) |
|---|
| 91 | temperature_range = ( |
|---|
| 92 | int(table['temp_min_k'][0]), |
|---|
| 93 | int(table['temp_max_k'][0]) |
|---|
| 94 | ) |
|---|
| 95 | pressure_range = ( |
|---|
| 96 | float(table['pressure_min_exponent_b'][0]), |
|---|
| 97 | float(table['pressure_max_exponent_b'][0]) |
|---|
| 98 | ) |
|---|
| 99 | return temperature_range, pressure_range |
|---|
| 100 | |
|---|
| 101 | |
|---|
| 102 | available_opacities = AvailableOpacities() |
|---|
| 103 | |
|---|
| 104 | |
|---|
| 105 | def get_atomic_database(): |
|---|
| 106 | db = Atom.query_database() |
|---|
| 107 | table = Table(db) |
|---|
| 108 | table.add_index('atom') |
|---|
| 109 | return table |
|---|
| 110 | |
|---|
| 111 | |
|---|
| 112 | def get_molecular_database(): |
|---|
| 113 | db = Molecule.query_database() |
|---|
| 114 | table = Table(db) |
|---|
| 115 | table.add_index('isotopologue') |
|---|
| 116 | return table |
|---|
| 117 | |
|---|
| 118 | class DownloadDace: |
|---|
| 119 | def __init__(self, |
|---|
| 120 | name='48Ti-16O', |
|---|
| 121 | line_list='Toto', |
|---|
| 122 | temperature_range=None, |
|---|
| 123 | pressure_range=None, |
|---|
| 124 | version=1, |
|---|
| 125 | output_dir="tmp" |
|---|
| 126 | ): |
|---|
| 127 | self.output_dir = output_dir |
|---|
| 128 | self.name = name |
|---|
| 129 | self.line_list = line_list |
|---|
| 130 | self.temperature_range = temperature_range |
|---|
| 131 | self.pressure_range = pressure_range |
|---|
| 132 | self.version = version |
|---|
| 133 | |
|---|
| 134 | def dace_download_molecule(self): |
|---|
| 135 | os.makedirs(self.output_dir, exist_ok=True) |
|---|
| 136 | archive_name = self.name + '__' + self.line_list + '.tar.gz' |
|---|
| 137 | path = os.path.join(self.output_dir, archive_name) |
|---|
| 138 | if os.path.exists(path): |
|---|
| 139 | print("Molecule already downloaded") |
|---|
| 140 | return path |
|---|
| 141 | Molecule.download( |
|---|
| 142 | self.name, |
|---|
| 143 | self.line_list, |
|---|
| 144 | float(self.version), |
|---|
| 145 | self.temperature_range, |
|---|
| 146 | self.pressure_range, |
|---|
| 147 | output_directory=self.output_dir, |
|---|
| 148 | output_filename=archive_name |
|---|
| 149 | ) |
|---|
| 150 | |
|---|
| 151 | return path |
|---|
| 152 | |
|---|
| 153 | |
|---|
| 154 | def dace_download_atom(self, charge=0): |
|---|
| 155 | os.makedirs(self.output_dir, exist_ok=True) |
|---|
| 156 | archive_name = self.name + '__' + self.line_list + '.tar.gz' |
|---|
| 157 | Atom.download( |
|---|
| 158 | self.name, charge, |
|---|
| 159 | self.line_list, float(self.version), |
|---|
| 160 | self.temperature_range, |
|---|
| 161 | self.pressure_range, |
|---|
| 162 | output_directory=self.output_dir, |
|---|
| 163 | output_filename=archive_name |
|---|
| 164 | ) |
|---|
| 165 | #### Atoms don't download anything for now, reason unknown |
|---|
| 166 | |
|---|
| 167 | return os.path.join(self.output_dir, archive_name) |
|---|
| 168 | |
|---|
| 169 | |
|---|
| 170 | def untar_bin_files(self, archive_name): |
|---|
| 171 | def bin_files(members): |
|---|
| 172 | for tarinfo in members: |
|---|
| 173 | if os.path.splitext(tarinfo.name)[1] == ".bin": |
|---|
| 174 | yield tarinfo |
|---|
| 175 | |
|---|
| 176 | path = os.path.join(self.output_dir, self.name + '__' + self.line_list) |
|---|
| 177 | with tarfile.open(archive_name, 'r') as tar: |
|---|
| 178 | tar.extractall(path=path , members=bin_files(tar)) |
|---|
| 179 | |
|---|
| 180 | return path |
|---|
| 181 | |
|---|
| 182 | |
|---|
| 183 | def opacity_dir_to_netcdf(self, opacity_dir, outpath): |
|---|
| 184 | outpath = self.output_dir + '/' + outpath |
|---|
| 185 | if os.path.exists(outpath): |
|---|
| 186 | print(f"{outpath} already exists") |
|---|
| 187 | return |
|---|
| 188 | temperature_grid = [] |
|---|
| 189 | pressure_grid = [] |
|---|
| 190 | wl_end = np.inf |
|---|
| 191 | |
|---|
| 192 | for dirpath, dirnames, filenames in os.walk(opacity_dir): |
|---|
| 193 | for filename in filenames: |
|---|
| 194 | if not filename.endswith('.bin'): |
|---|
| 195 | continue |
|---|
| 196 | |
|---|
| 197 | # Wavenumber points from range given in the file names |
|---|
| 198 | temperature = int(filename.split('_')[3]) |
|---|
| 199 | sign = 1 if filename.split('_')[4][0] == 'p' else -1 |
|---|
| 200 | pressure = 10 ** (sign * float(filename.split('_')[4][1:].split('.')[0]) / 100) |
|---|
| 201 | |
|---|
| 202 | wl_start = int(filename.split('_')[1]) |
|---|
| 203 | wl_end = min(wl_end, int(filename.split('_')[2])) |
|---|
| 204 | wlen = np.arange(wl_start, wl_end, 0.01) |
|---|
| 205 | |
|---|
| 206 | # catch divide by zero warning here: |
|---|
| 207 | with warnings.catch_warnings(): |
|---|
| 208 | warnings.simplefilter('ignore', RuntimeWarning) |
|---|
| 209 | |
|---|
| 210 | # Convert to micron |
|---|
| 211 | wavelength = 1 / wlen / 1e-4 |
|---|
| 212 | |
|---|
| 213 | unique_wavelengths = wavelength[1:][::-1] |
|---|
| 214 | temperature_grid.append(temperature) |
|---|
| 215 | pressure_grid.append(pressure) |
|---|
| 216 | |
|---|
| 217 | tgrid = np.sort(list(set(temperature_grid))) |
|---|
| 218 | pgrid = np.sort(list(set(pressure_grid))) |
|---|
| 219 | |
|---|
| 220 | if len(pgrid) == 1: |
|---|
| 221 | extrapolate_pgrid = True |
|---|
| 222 | pgrid = np.concatenate([pgrid, 10 ** (-1 * np.log10(pgrid))]) |
|---|
| 223 | else: |
|---|
| 224 | extrapolate_pgrid = False |
|---|
| 225 | |
|---|
| 226 | opacity_grid = np.zeros( |
|---|
| 227 | (len(tgrid), len(pgrid), len(unique_wavelengths)), dtype='float32' |
|---|
| 228 | ) |
|---|
| 229 | |
|---|
| 230 | for dirpath, dirnames, filenames in os.walk(opacity_dir): |
|---|
| 231 | for filename in filenames: |
|---|
| 232 | if not filename.endswith('.bin'): |
|---|
| 233 | continue |
|---|
| 234 | |
|---|
| 235 | opacity = np.fromfile( |
|---|
| 236 | os.path.join(dirpath, filename), dtype=np.float32 |
|---|
| 237 | )[1:][::-1] |
|---|
| 238 | |
|---|
| 239 | # Wavenumber points from range given in the file names |
|---|
| 240 | temperature = int(filename.split('_')[3]) |
|---|
| 241 | sign = 1 if filename.split('_')[4][0] == 'p' else -1 |
|---|
| 242 | pressure = 10 ** (sign * float(filename.split('_')[4][1:].split('.')[0]) / 100) |
|---|
| 243 | |
|---|
| 244 | temperature_ind = np.argmin(np.abs(tgrid - temperature)) |
|---|
| 245 | pressure_ind = np.argmin(np.abs(pgrid - pressure)) |
|---|
| 246 | |
|---|
| 247 | opacity_grid[temperature_ind, pressure_ind, :] = opacity[:len(unique_wavelengths)] |
|---|
| 248 | |
|---|
| 249 | if extrapolate_pgrid: |
|---|
| 250 | for dirpath, dirnames, filenames in os.walk(opacity_dir): |
|---|
| 251 | for filename in filenames: |
|---|
| 252 | opacity = np.fromfile( |
|---|
| 253 | os.path.join(dirpath, filename), dtype=np.float32 |
|---|
| 254 | )[1:][::-1] |
|---|
| 255 | |
|---|
| 256 | # Wavenumber points from range given in the file names |
|---|
| 257 | temperature = int(filename.split('_')[3]) |
|---|
| 258 | # *Flip the sign for the extrapolated grid point in pressure* |
|---|
| 259 | sign = -1 if filename.split('_')[4][0] == 'p' else 1 |
|---|
| 260 | pressure = 10 ** (sign * float(filename.split('_')[4][1:].split('.')[0]) / 100) |
|---|
| 261 | |
|---|
| 262 | temperature_ind = np.argmin(np.abs(tgrid - temperature)) |
|---|
| 263 | pressure_ind = np.argmin(np.abs(pgrid - pressure)) |
|---|
| 264 | |
|---|
| 265 | opacity_grid[temperature_ind, pressure_ind, :] = opacity |
|---|
| 266 | |
|---|
| 267 | ds = xr.Dataset( |
|---|
| 268 | data_vars=dict( |
|---|
| 269 | opacity=(["temperature", "pressure", "wavelength"], |
|---|
| 270 | opacity_grid) |
|---|
| 271 | ), |
|---|
| 272 | coords=dict( |
|---|
| 273 | temperature=(["temperature"], tgrid), |
|---|
| 274 | pressure=(["pressure"], pgrid), |
|---|
| 275 | wavelength=unique_wavelengths |
|---|
| 276 | ) |
|---|
| 277 | ) |
|---|
| 278 | if not os.path.exists(os.path.dirname(outpath)): |
|---|
| 279 | os.makedirs(os.path.dirname(outpath), exist_ok=True) |
|---|
| 280 | |
|---|
| 281 | ds.to_netcdf(outpath if outpath.endswith(".nc") else outpath + '.nc', |
|---|
| 282 | encoding={'opacity': {'dtype': 'float32'}}) |
|---|
| 283 | |
|---|
| 284 | |
|---|
| 285 | def clean_up(self, bin_dir, archive_name): |
|---|
| 286 | os.remove(archive_name) |
|---|
| 287 | shutil.rmtree(bin_dir) |
|---|
| 288 | |
|---|
| 289 | |
|---|
| 290 | def download_molecule( |
|---|
| 291 | molecule=None, |
|---|
| 292 | isotopologue=None, |
|---|
| 293 | line_list='first-found', |
|---|
| 294 | temperature_range=None, |
|---|
| 295 | pressure_range=None, |
|---|
| 296 | version=None, |
|---|
| 297 | output_dir="datadir/dace" |
|---|
| 298 | ): |
|---|
| 299 | """ |
|---|
| 300 | Download molecular opacity data from DACE. |
|---|
| 301 | |
|---|
| 302 | .. warning:: |
|---|
| 303 | This generates *very* large files. Only run this |
|---|
| 304 | method if you have ~6 GB available per molecule. |
|---|
| 305 | |
|---|
| 306 | Parameters |
|---|
| 307 | ---------- |
|---|
| 308 | isotopologue : str |
|---|
| 309 | For example, "1H2-16O" for water. |
|---|
| 310 | molecule : str |
|---|
| 311 | Common name for the molecule, for example: "H2O" |
|---|
| 312 | line_list : str, default is ``'first-found'``, optional |
|---|
| 313 | For example, "POKAZATEL" for water. By default, the first available |
|---|
| 314 | line list for this isotopologue is chosen. |
|---|
| 315 | temperature_range : tuple, optional |
|---|
| 316 | Tuple of integers specifying the min and max |
|---|
| 317 | temperature requested. Defaults to the full |
|---|
| 318 | range of available temperatures. |
|---|
| 319 | pressure_range : tuple, optional |
|---|
| 320 | Tuple of floags specifying the log base 10 of the |
|---|
| 321 | min and max pressure [bar] requested. Defaults to the full |
|---|
| 322 | range of available pressures. |
|---|
| 323 | version : float, optional |
|---|
| 324 | Version number of the line list in DACE. Defaults to the |
|---|
| 325 | latest version. |
|---|
| 326 | """ |
|---|
| 327 | if molecule is not None: |
|---|
| 328 | isotopologue = species_name_to_common_isotopologue_name(molecule) |
|---|
| 329 | |
|---|
| 330 | available_line_lists = available_opacities.get_molecular_line_lists(isotopologue) |
|---|
| 331 | print(f"Available line lists for {isotopologue}:") |
|---|
| 332 | print(available_line_lists) |
|---|
| 333 | |
|---|
| 334 | if line_list is None: |
|---|
| 335 | return |
|---|
| 336 | |
|---|
| 337 | if line_list == 'first-found': |
|---|
| 338 | line_list = sorted(list(available_line_lists)).pop() |
|---|
| 339 | logging.warning(f"Using first-found line list for {isotopologue}: '{line_list}'") |
|---|
| 340 | |
|---|
| 341 | elif line_list not in available_line_lists: |
|---|
| 342 | raise ValueError(f"The requested '{line_list}' is not in the set of " |
|---|
| 343 | f"available line lists {available_line_lists}.") |
|---|
| 344 | |
|---|
| 345 | if version is None: |
|---|
| 346 | version = available_opacities.get_molecular_latest_version(isotopologue, line_list) |
|---|
| 347 | logging.warning(f"Using latest version of the line " |
|---|
| 348 | f"list '{line_list}' for {isotopologue}: {version}") |
|---|
| 349 | |
|---|
| 350 | if temperature_range is None or pressure_range is None: |
|---|
| 351 | dace_temp_range, dace_press_range = available_opacities.get_molecular_pT_range( |
|---|
| 352 | isotopologue, line_list |
|---|
| 353 | ) |
|---|
| 354 | |
|---|
| 355 | if temperature_range is None: |
|---|
| 356 | temperature_range = dace_temp_range |
|---|
| 357 | |
|---|
| 358 | if pressure_range is None: |
|---|
| 359 | pressure_range = dace_press_range |
|---|
| 360 | |
|---|
| 361 | download_dace = DownloadDace(isotopologue, line_list, temperature_range, pressure_range, version, output_dir) |
|---|
| 362 | archive_name = download_dace.dace_download_molecule() |
|---|
| 363 | |
|---|
| 364 | bin_dir = download_dace.untar_bin_files(archive_name) |
|---|
| 365 | |
|---|
| 366 | download_dace.opacity_dir_to_netcdf(bin_dir, isotopologue + '__' + line_list + '.nc') |
|---|
| 367 | # clean_up(bin_dir, archive_name) |
|---|
| 368 | return isotopologue + '__' + line_list + '.nc' |
|---|
| 369 | |
|---|
| 370 | |
|---|
| 371 | def download_atom(atom, charge, line_list='first-found', |
|---|
| 372 | temperature_range=None, pressure_range=None, version=None, output_dir="tmp"): |
|---|
| 373 | """ |
|---|
| 374 | Download atomic opacity data from DACE. |
|---|
| 375 | |
|---|
| 376 | .. warning:: |
|---|
| 377 | This generates *very* large files. Only run this |
|---|
| 378 | method if you have ~6 GB available per molecule. |
|---|
| 379 | |
|---|
| 380 | Parameters |
|---|
| 381 | ---------- |
|---|
| 382 | atom : str |
|---|
| 383 | For example, "Na" for sodium. |
|---|
| 384 | charge : int |
|---|
| 385 | For example, 0 for neutral. |
|---|
| 386 | line_list : str, default is ``'first-found'``, optional |
|---|
| 387 | For example, "Kurucz". By default, the first available |
|---|
| 388 | line list for this atom/charge is chosen. |
|---|
| 389 | temperature_range : tuple, optional |
|---|
| 390 | Tuple of integers specifying the min and max |
|---|
| 391 | temperature requested. Defaults to the full |
|---|
| 392 | range of available temperatures. |
|---|
| 393 | pressure_range : tuple, optional |
|---|
| 394 | Tuple of floags specifying the log base 10 of the |
|---|
| 395 | min and max pressure [bar] requested. Defaults to the full |
|---|
| 396 | range of available pressures. |
|---|
| 397 | version : float, optional |
|---|
| 398 | Version number of the line list in DACE. Defaults to the |
|---|
| 399 | latest version. |
|---|
| 400 | """ |
|---|
| 401 | available_line_lists = available_opacities.get_atomic_line_lists(atom) |
|---|
| 402 | print(f"Available line lists for {atom}:") |
|---|
| 403 | print(available_line_lists) |
|---|
| 404 | |
|---|
| 405 | if line_list is None: |
|---|
| 406 | return |
|---|
| 407 | |
|---|
| 408 | if line_list == 'first-found': |
|---|
| 409 | line_list = sorted(list(available_line_lists)).pop() |
|---|
| 410 | logging.warning(f"Using first-found line list for {atom}: '{line_list}'") |
|---|
| 411 | |
|---|
| 412 | elif line_list not in available_line_lists: |
|---|
| 413 | raise ValueError(f"The requested '{line_list}' is not in the set of " |
|---|
| 414 | f"available line lists {available_line_lists}.") |
|---|
| 415 | |
|---|
| 416 | if temperature_range is None or pressure_range is None: |
|---|
| 417 | dace_temp_range, dace_press_range = available_opacities.get_atomic_pT_range( |
|---|
| 418 | atom, charge, line_list |
|---|
| 419 | ) |
|---|
| 420 | |
|---|
| 421 | if version is None: |
|---|
| 422 | version = available_opacities.get_atomic_latest_version(atom, charge, line_list) |
|---|
| 423 | logging.warning(f"Using latest version of the line " |
|---|
| 424 | f"list '{line_list}' for {atom}: {version}") |
|---|
| 425 | |
|---|
| 426 | if temperature_range is None: |
|---|
| 427 | temperature_range = dace_temp_range |
|---|
| 428 | |
|---|
| 429 | if pressure_range is None: |
|---|
| 430 | pressure_range = dace_press_range |
|---|
| 431 | |
|---|
| 432 | download_dace = DownloadDace(atom, line_list, temperature_range, pressure_range, version, output_dir) |
|---|
| 433 | |
|---|
| 434 | archive_name = download_dace.dace_download_atom(charge) |
|---|
| 435 | bin_dir = download_dace.untar_bin_files(archive_name) |
|---|
| 436 | |
|---|
| 437 | download_dace.opacity_dir_to_netcdf(bin_dir, atom + '_' + str(int(charge)) + '__' + line_list + '.nc') |
|---|
| 438 | # download_dace.clean_up(bin_dir, archive_name) |
|---|
