API tools faq
paste
Advertisement
mabruzzo

alternative grackle_wrapper.pyx

mabruzzo
Dec 12th, 2022
81
0
Never
Add comment
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
text 39.18 KB | None | 0 0
raw download clone embed print report
  1. ########################################################################
  2. #
  3. # Python Grackle wrapper
  4. #
  5. #
  6. # Copyright (c) 2013, Enzo/Grackle Development Team.
  7. #
  8. # Distributed under the terms of the Enzo Public Licence.
  9. #
  10. # The full license is in the file LICENSE, distributed with this
  11. # software.
  12. ########################################################################
  13.  
  14. from pygrackle.utilities.physical_constants import \
  15. boltzmann_constant_cgs, \
  16. mass_hydrogen_cgs
  17.  
  18. from .grackle_defs cimport *
  19. import numpy as np
  20. cimport numpy as np
  21.  
  22. cdef void _fetch_struct_ptrs(object chem_data, c_chemistry_data **chem_ptr,
  23. c_chemistry_data_storage **rates_ptr,
  24. c_code_units **units_ptr) except *:
  25. # annotating this with `except *` is essential for proper exception
  26. # handling if one of the implicit type checks fails
  27. cdef _wrapped_c_chemistry_data chemistry_wrapper = chem_data.data
  28. chem_ptr[0] = &chemistry_wrapper.data
  29.  
  30. cdef _chemistry_storage_and_units other_wrapper = chem_data.other
  31. rates_ptr[0] = &other_wrapper.rates
  32. units_ptr[0] = &other_wrapper.units
  33.  
  34. class chemistry_data:
  35. # this is not an extension type to simplify __setattr__
  36.  
  37. def __init__(self):
  38. self.data = _wrapped_c_chemistry_data()
  39. self.other = _chemistry_storage_and_units()
  40.  
  41. def __getattribute__(self, name):
  42. # This has been implemented to maintain backwards compatibility.
  43.  
  44. # first check if the name is contained within self.data or self.other
  45. try:
  46. return self.data[name]
  47. except:
  48. pass
  49.  
  50. try:
  51. return getattr(self.other, name)
  52. except:
  53. pass
  54.  
  55. # fall back to the default implementation
  56. return super().__getattribute__(name)
  57.  
  58. def __setattr__(self, name, value):
  59. # This has been implemented to maintain backwards compatibility.
  60.  
  61. # first check if the name is contained within self.data or self.other
  62. try:
  63. self.data[name] = value
  64. return
  65. except KeyError: # raise other Exceptions
  66. pass
  67.  
  68. try:
  69. setattr(self, name, value)
  70. except AttributeError: # raise other Exceptions
  71. pass
  72.  
  73. # fall back to the default implementation
  74. return super().__setattr__(name,value)
  75.  
  76. def initialize(self):
  77. cdef c_chemistry_data *chem_ptr = NULL
  78. cdef c_chemistry_data_storage *rates_ptr = NULL
  79. cdef c_code_units *units_ptr = NULL
  80. _fetch_struct_ptrs(self, &chem_ptr, &rates_ptr, &units_ptr)
  81.  
  82. ret = _initialize_chemistry_data(chem_ptr, rates_ptr, units_ptr)
  83. if ret is None:
  84. raise RuntimeError("Error initializing chemistry")
  85. return ret
  86.  
  87. def set_velocity_units(self):
  88. self.other.set_velocity_units()
  89.  
  90. def get_velocity_units(self):
  91. return self.other.get_velocity_units()
  92.  
  93. cdef class _chemistry_storage_and_units:
  94. cdef c_chemistry_data_storage rates
  95. cdef c_code_units units
  96.  
  97. def set_velocity_units(self):
  98. set_velocity_units(&self.units)
  99.  
  100. def get_velocity_units(self):
  101. return get_velocity_units(&self.units)
  102.  
  103. property k1:
  104. def __get__(self):
  105. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.k1)
  106. return np.asarray(memview)
  107.  
  108. property k2:
  109. def __get__(self):
  110. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.k2)
  111. return np.asarray(memview)
  112.  
  113. property k3:
  114. def __get__(self):
  115. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.k3)
  116. return np.asarray(memview)
  117.  
  118. property k4:
  119. def __get__(self):
  120. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.k4)
  121. return np.asarray(memview)
  122.  
  123. property k5:
  124. def __get__(self):
  125. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.k5)
  126. return np.asarray(memview)
  127.  
  128. property k6:
  129. def __get__(self):
  130. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.k6)
  131. return np.asarray(memview)
  132.  
  133. property k7:
  134. def __get__(self):
  135. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.k7)
  136. return np.asarray(memview)
  137.  
  138. property k8:
  139. def __get__(self):
  140. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.k8)
  141. return np.asarray(memview)
  142.  
  143. property k9:
  144. def __get__(self):
  145. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.k9)
  146. return np.asarray(memview)
  147.  
  148. property k10:
  149. def __get__(self):
  150. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.k10)
  151. return np.asarray(memview)
  152.  
  153. property k11:
  154. def __get__(self):
  155. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.k11)
  156. return np.asarray(memview)
  157.  
  158. property k12:
  159. def __get__(self):
  160. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.k12)
  161. return np.asarray(memview)
  162.  
  163. property k13:
  164. def __get__(self):
  165. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.k13)
  166. return np.asarray(memview)
  167.  
  168. property k14:
  169. def __get__(self):
  170. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.k14)
  171. return np.asarray(memview)
  172.  
  173. property k15:
  174. def __get__(self):
  175. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.k15)
  176. return np.asarray(memview)
  177.  
  178. property k16:
  179. def __get__(self):
  180. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.k16)
  181. return np.asarray(memview)
  182.  
  183. property k17:
  184. def __get__(self):
  185. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.k17)
  186. return np.asarray(memview)
  187.  
  188. property k18:
  189. def __get__(self):
  190. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.k18)
  191. return np.asarray(memview)
  192.  
  193. property k19:
  194. def __get__(self):
  195. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.k19)
  196. return np.asarray(memview)
  197.  
  198. property k20:
  199. def __get__(self):
  200. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.k20)
  201. return np.asarray(memview)
  202.  
  203. property k21:
  204. def __get__(self):
  205. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.k21)
  206. return np.asarray(memview)
  207.  
  208. property k22:
  209. def __get__(self):
  210. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.k22)
  211. return np.asarray(memview)
  212.  
  213. property k23:
  214. def __get__(self):
  215. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.k23)
  216. return np.asarray(memview)
  217.  
  218. property k13dd:
  219. def __get__(self):
  220. cdef double[:] memview = <double[:self.NumberOfTemperatureBins*14]>(<double*> self.rates.k13dd)
  221. return np.asarray(memview)
  222.  
  223. property k24:
  224. def __get__(self):
  225. return self.rates.k24
  226. def __set__(self, val):
  227. self.rates.k24 = val
  228.  
  229. property k25:
  230. def __get__(self):
  231. return self.rates.k25
  232. def __set__(self, val):
  233. self.rates.k25 = val
  234.  
  235. property k26:
  236. def __get__(self):
  237. return self.rates.k26
  238. def __set__(self, val):
  239. self.rates.k26 = val
  240.  
  241. property k27:
  242. def __get__(self):
  243. return self.rates.k27
  244. def __set__(self, val):
  245. self.rates.k27 = val
  246.  
  247. property k28:
  248. def __get__(self):
  249. return self.rates.k28
  250. def __set__(self, val):
  251. self.rates.k28 = val
  252.  
  253. property k29:
  254. def __get__(self):
  255. return self.rates.k29
  256. def __set__(self, val):
  257. self.rates.k29 = val
  258.  
  259. property k30:
  260. def __get__(self):
  261. return self.rates.k30
  262. def __set__(self, val):
  263. self.rates.k30 = val
  264.  
  265. property k31:
  266. def __get__(self):
  267. return self.rates.k31
  268. def __set__(self, val):
  269. self.rates.k31 = val
  270.  
  271. property k50:
  272. def __get__(self):
  273. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.k50)
  274. return np.asarray(memview)
  275.  
  276. property k51:
  277. def __get__(self):
  278. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.k51)
  279. return np.asarray(memview)
  280.  
  281. property k52:
  282. def __get__(self):
  283. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.k52)
  284. return np.asarray(memview)
  285.  
  286. property k53:
  287. def __get__(self):
  288. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.k53)
  289. return np.asarray(memview)
  290.  
  291. property k54:
  292. def __get__(self):
  293. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.k54)
  294. return np.asarray(memview)
  295.  
  296. property k55:
  297. def __get__(self):
  298. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.k55)
  299. return np.asarray(memview)
  300.  
  301. property k56:
  302. def __get__(self):
  303. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.k56)
  304. return np.asarray(memview)
  305.  
  306. property k57:
  307. def __get__(self):
  308. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.k57)
  309. return np.asarray(memview)
  310.  
  311. property k58:
  312. def __get__(self):
  313. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.k58)
  314. return np.asarray(memview)
  315.  
  316. property h2dust:
  317. def __get__(self):
  318. cdef double[:] memview = <double[:self.NumberOfTemperatureBins*self.NumberOfDustTemperatureBins]>(<double*> self.rates.h2dust)
  319. return np.asarray(memview)
  320.  
  321. property n_cr_n:
  322. def __get__(self):
  323. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.n_cr_n)
  324. return np.asarray(memview)
  325.  
  326. property n_cr_d1:
  327. def __get__(self):
  328. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.n_cr_d1)
  329. return np.asarray(memview)
  330.  
  331. property n_cr_d2:
  332. def __get__(self):
  333. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.n_cr_d2)
  334. return np.asarray(memview)
  335.  
  336. property ceHI:
  337. def __get__(self):
  338. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.ceHI)
  339. return np.asarray(memview)
  340.  
  341. property ceHeI:
  342. def __get__(self):
  343. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.ceHeI)
  344. return np.asarray(memview)
  345.  
  346. property ceHeII:
  347. def __get__(self):
  348. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.ceHeII)
  349. return np.asarray(memview)
  350.  
  351. property ciHI:
  352. def __get__(self):
  353. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.ciHI)
  354. return np.asarray(memview)
  355.  
  356. property ciHeI:
  357. def __get__(self):
  358. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.ciHeI)
  359. return np.asarray(memview)
  360.  
  361. property ciHeIS:
  362. def __get__(self):
  363. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.ciHeIS)
  364. return np.asarray(memview)
  365.  
  366. property ciHeII:
  367. def __get__(self):
  368. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.ciHeII)
  369. return np.asarray(memview)
  370.  
  371. property reHII:
  372. def __get__(self):
  373. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.reHII)
  374. return np.asarray(memview)
  375.  
  376. property reHeII1:
  377. def __get__(self):
  378. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.reHeII1)
  379. return np.asarray(memview)
  380.  
  381. property reHeII2:
  382. def __get__(self):
  383. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.reHeII2)
  384. return np.asarray(memview)
  385.  
  386. property reHeIII:
  387. def __get__(self):
  388. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.reHeIII)
  389. return np.asarray(memview)
  390.  
  391. property brem:
  392. def __get__(self):
  393. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.brem)
  394. return np.asarray(memview)
  395.  
  396. property comp:
  397. def __get__(self):
  398. return self.rates.comp
  399. def __set__(self, val):
  400. self.rates.comp = val
  401.  
  402. property comp_xray:
  403. def __get__(self):
  404. return self.rates.comp_xray
  405. def __set__(self, val):
  406. self.rates.comp_xray = val
  407.  
  408. property temp_xray:
  409. def __get__(self):
  410. return self.rates.temp_xray
  411. def __set__(self, val):
  412. self.rates.temp_xray = val
  413.  
  414. property piHI:
  415. def __get__(self):
  416. return self.rates.piHI
  417. def __set__(self, val):
  418. self.rates.piHI = val
  419.  
  420. property piHeI:
  421. def __get__(self):
  422. return self.rates.piHeI
  423. def __set__(self, val):
  424. self.rates.piHeI = val
  425.  
  426. property piHeII:
  427. def __get__(self):
  428. return self.rates.piHeII
  429. def __set__(self, val):
  430. self.rates.piHeII = val
  431.  
  432. property crsHI:
  433. def __get__(self):
  434. return self.rates.crsHI
  435. def __set__(self, val):
  436. self.rates.crsHI = val
  437.  
  438. property crsHeI:
  439. def __get__(self):
  440. return self.rates.crsHeI
  441. def __set__(self, val):
  442. self.rates.crsHeI = val
  443.  
  444. property crsHeII:
  445. def __get__(self):
  446. return self.rates.crsHeII
  447. def __set__(self, val):
  448. self.rates.crsHeII = val
  449.  
  450. property hyd01k:
  451. def __get__(self):
  452. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.hyd01k)
  453. return np.asarray(memview)
  454.  
  455. property h2k01:
  456. def __get__(self):
  457. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.h2k01)
  458. return np.asarray(memview)
  459.  
  460. property vibh:
  461. def __get__(self):
  462. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.vibh)
  463. return np.asarray(memview)
  464.  
  465. property roth:
  466. def __get__(self):
  467. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.roth)
  468. return np.asarray(memview)
  469.  
  470. property rotl:
  471. def __get__(self):
  472. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.rotl)
  473. return np.asarray(memview)
  474.  
  475. property GP99LowDensityLimit:
  476. def __get__(self):
  477. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.GP99LowDensityLimit)
  478. return np.asarray(memview)
  479.  
  480. property GP99HighDensityLimit:
  481. def __get__(self):
  482. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.GP99HighDensityLimit)
  483. return np.asarray(memview)
  484.  
  485. property GAHI:
  486. def __get__(self):
  487. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.GAHI)
  488. return np.asarray(memview)
  489.  
  490. property GAH2:
  491. def __get__(self):
  492. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.GAH2)
  493. return np.asarray(memview)
  494.  
  495. property GAHe:
  496. def __get__(self):
  497. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.GAHe)
  498. return np.asarray(memview)
  499.  
  500. property GAHp:
  501. def __get__(self):
  502. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.GAHp)
  503. return np.asarray(memview)
  504.  
  505. property GAel:
  506. def __get__(self):
  507. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.GAel)
  508. return np.asarray(memview)
  509.  
  510. property H2LTE:
  511. def __get__(self):
  512. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.H2LTE)
  513. return np.asarray(memview)
  514.  
  515. property HDlte:
  516. def __get__(self):
  517. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.HDlte)
  518. return np.asarray(memview)
  519.  
  520. property HDlow:
  521. def __get__(self):
  522. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.HDlow)
  523. return np.asarray(memview)
  524.  
  525. property cieco:
  526. def __get__(self):
  527. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.cieco)
  528. return np.asarray(memview)
  529.  
  530. property gammah:
  531. def __get__(self):
  532. return self.rates.gammah
  533. def __set__(self, val):
  534. self.rates.gammah = val
  535.  
  536. property regr:
  537. def __get__(self):
  538. if not self.dust_chemistry and not self.h2_on_dust:
  539. return 0
  540. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.regr)
  541. return np.asarray(memview)
  542.  
  543. property gamma_isrf:
  544. def __get__(self):
  545. return self.rates.gamma_isrf
  546. def __set__(self, val):
  547. self.rates.gamma_isrf = val
  548.  
  549. property gas_grain:
  550. def __get__(self):
  551. if not self.dust_chemistry and not self.h2_on_dust:
  552. return 0
  553. cdef double[:] memview = <double[:self.NumberOfTemperatureBins]>(<double*> self.rates.gas_grain)
  554. return np.asarray(memview)
  555.  
  556. property comoving_coordinates:
  557. def __get__(self):
  558. return self.units.comoving_coordinates
  559. def __set__(self, val):
  560. self.units.comoving_coordinates = val
  561.  
  562. property density_units:
  563. def __get__(self):
  564. return self.units.density_units
  565. def __set__(self, val):
  566. self.units.density_units = val
  567.  
  568. property length_units:
  569. def __get__(self):
  570. return self.units.length_units
  571. def __set__(self, val):
  572. self.units.length_units = val
  573.  
  574. property time_units:
  575. def __get__(self):
  576. return self.units.time_units
  577. def __set__(self, val):
  578. self.units.time_units = val
  579.  
  580. property velocity_units:
  581. def __get__(self):
  582. return self.units.velocity_units
  583. def __set__(self, val):
  584. self.units.velocity_units = val
  585.  
  586. property a_units:
  587. def __get__(self):
  588. return self.units.a_units
  589. def __set__(self, val):
  590. self.units.a_units = val
  591.  
  592. property a_value:
  593. def __get__(self):
  594. return self.units.a_value
  595. def __set__(self, val):
  596. self.units.a_value = val
  597.  
  598. property temperature_units:
  599. def __get__(self):
  600. return get_temperature_units(&self.units)
  601.  
  602. property cooling_units:
  603. def __get__(self):
  604. tbase1 = self.time_units
  605. if self.comoving_coordinates:
  606. xbase1 = self.length_units / \
  607. (self.a_value * self.a_units)
  608. dbase1 = self.density_units * \
  609. (self.a_value * self.a_units)**3
  610. else:
  611. xbase1 = self.length_units / \
  612. self.a_units
  613. dbase1 = self.density_units * \
  614. self.a_units**3
  615.  
  616. coolunit = (self.a_units**5 * xbase1**2 *
  617. mass_hydrogen_cgs**2) / (tbase1**3 * dbase1)
  618. return coolunit
  619.  
  620. property energy_units:
  621. def __get__(self):
  622. return self.velocity_units**2
  623.  
  624. property pressure_units:
  625. def __get__(self):
  626. return self.density_units * self.energy_units
  627.  
  628. cdef gr_float* get_field(fc, name):
  629. cdef np.ndarray rv = fc.get(name, None)
  630. if rv is None:
  631. return NULL
  632. else:
  633. return <gr_float *> rv.data
  634.  
  635. def solve_chemistry(fc, my_dt):
  636. cdef double dt_value = <double> my_dt
  637.  
  638. cdef c_chemistry_data *chem_ptr = NULL
  639. cdef c_chemistry_data_storage *rates_ptr = NULL
  640. cdef c_code_units *units_ptr = NULL
  641. _fetch_struct_ptrs(fc.chemistry_data, &chem_ptr, &rates_ptr, &units_ptr)
  642.  
  643. cdef int grid_dimension
  644. grid_dimension = fc["density"].shape[0]
  645. cdef np.ndarray ref_gs, ref_ge
  646. ref_gs = np.zeros(3, dtype="int32")
  647. ref_ge = np.zeros(3, dtype="int32")
  648. ref_ge[0] = grid_dimension -1
  649.  
  650. cdef c_field_data my_fields
  651. my_fields.grid_rank = 1
  652. my_fields.grid_dimension = &grid_dimension
  653. my_fields.grid_start = <int *> ref_gs.data
  654. my_fields.grid_end = <int *> ref_ge.data
  655. my_fields.density = get_field(fc, "density")
  656. my_fields.internal_energy = get_field(fc, "energy")
  657. my_fields.HI_density = get_field(fc, "HI")
  658. my_fields.HII_density = get_field(fc, "HII")
  659. my_fields.HM_density = get_field(fc, "HM")
  660. my_fields.HeI_density = get_field(fc, "HeI")
  661. my_fields.HeII_density = get_field(fc, "HeII")
  662. my_fields.HeIII_density = get_field(fc, "HeIII")
  663. my_fields.H2I_density = get_field(fc, "H2I")
  664. my_fields.H2II_density = get_field(fc, "H2II")
  665. my_fields.DI_density = get_field(fc, "DI")
  666. my_fields.DII_density = get_field(fc, "DII")
  667. my_fields.HDI_density = get_field(fc, "HDI")
  668. my_fields.e_density = get_field(fc, "de")
  669. my_fields.metal_density = get_field(fc, "metal")
  670. my_fields.dust_density = get_field(fc, "dust")
  671. my_fields.RT_heating_rate = get_field(fc, "RT_heating_rate")
  672. my_fields.volumetric_heating_rate = get_field(fc, "volumetric_heating_rate")
  673. my_fields.specific_heating_rate = get_field(fc, "specific_heating_rate")
  674.  
  675. c_local_solve_chemistry(
  676. chem_ptr,
  677. rates_ptr,
  678. units_ptr,
  679. &my_fields,
  680. my_dt)
  681.  
  682. def calculate_cooling_time(fc):
  683. cdef c_chemistry_data *chem_ptr = NULL
  684. cdef c_chemistry_data_storage *rates_ptr = NULL
  685. cdef c_code_units *units_ptr = NULL
  686. _fetch_struct_ptrs(fc.chemistry_data, &chem_ptr, &rates_ptr, &units_ptr)
  687.  
  688. cdef int grid_dimension
  689. grid_dimension = fc["density"].shape[0]
  690. cdef np.ndarray ref_gs, ref_ge
  691. ref_gs = np.zeros(3, dtype="int32")
  692. ref_ge = np.zeros(3, dtype="int32")
  693. ref_ge[0] = grid_dimension -1
  694.  
  695. cdef c_field_data my_fields
  696. my_fields.grid_rank = 1
  697. my_fields.grid_dimension = &grid_dimension
  698. my_fields.grid_start = <int *> ref_gs.data
  699. my_fields.grid_end = <int *> ref_ge.data
  700. my_fields.density = get_field(fc, "density")
  701. my_fields.internal_energy = get_field(fc, "energy")
  702. my_fields.x_velocity = get_field(fc, "x-velocity")
  703. my_fields.y_velocity = get_field(fc, "y-velocity")
  704. my_fields.z_velocity = get_field(fc, "z-velocity")
  705. my_fields.HI_density = get_field(fc, "HI")
  706. my_fields.HII_density = get_field(fc, "HII")
  707. my_fields.HM_density = get_field(fc, "HM")
  708. my_fields.HeI_density = get_field(fc, "HeI")
  709. my_fields.HeII_density = get_field(fc, "HeII")
  710. my_fields.HeIII_density = get_field(fc, "HeIII")
  711. my_fields.H2I_density = get_field(fc, "H2I")
  712. my_fields.H2II_density = get_field(fc, "H2II")
  713. my_fields.DI_density = get_field(fc, "DI")
  714. my_fields.DII_density = get_field(fc, "DII")
  715. my_fields.HDI_density = get_field(fc, "HDI")
  716. my_fields.e_density = get_field(fc, "de")
  717. my_fields.metal_density = get_field(fc, "metal")
  718. my_fields.dust_density = get_field(fc, "dust")
  719. my_fields.RT_heating_rate = get_field(fc, "RT_heating_rate")
  720. my_fields.volumetric_heating_rate = get_field(fc, "volumetric_heating_rate")
  721. my_fields.specific_heating_rate = get_field(fc, "specific_heating_rate")
  722. cdef gr_float *cooling_time = get_field(fc, "cooling_time")
  723.  
  724. c_local_calculate_cooling_time(
  725. chem_ptr,
  726. rates_ptr,
  727. units_ptr,
  728. &my_fields,
  729. cooling_time)
  730.  
  731. def calculate_gamma(fc):
  732. cdef c_chemistry_data *chem_ptr = NULL
  733. cdef c_chemistry_data_storage *rates_ptr = NULL
  734. cdef c_code_units *units_ptr = NULL
  735. _fetch_struct_ptrs(fc.chemistry_data, &chem_ptr, &rates_ptr, &units_ptr)
  736.  
  737. cdef int grid_dimension
  738. grid_dimension = fc["density"].shape[0]
  739. cdef np.ndarray ref_gs, ref_ge
  740. ref_gs = np.zeros(3, dtype="int32")
  741. ref_ge = np.zeros(3, dtype="int32")
  742. ref_ge[0] = grid_dimension -1
  743.  
  744. cdef c_field_data my_fields
  745. my_fields.grid_rank = 1
  746. my_fields.grid_dimension = &grid_dimension
  747. my_fields.grid_start = <int *> ref_gs.data
  748. my_fields.grid_end = <int *> ref_ge.data
  749. my_fields.density = get_field(fc, "density")
  750. my_fields.internal_energy = get_field(fc, "energy")
  751. my_fields.x_velocity = get_field(fc, "x-velocity")
  752. my_fields.y_velocity = get_field(fc, "y-velocity")
  753. my_fields.z_velocity = get_field(fc, "z-velocity")
  754. my_fields.HI_density = get_field(fc, "HI")
  755. my_fields.HII_density = get_field(fc, "HII")
  756. my_fields.HM_density = get_field(fc, "HM")
  757. my_fields.HeI_density = get_field(fc, "HeI")
  758. my_fields.HeII_density = get_field(fc, "HeII")
  759. my_fields.HeIII_density = get_field(fc, "HeIII")
  760. my_fields.H2I_density = get_field(fc, "H2I")
  761. my_fields.H2II_density = get_field(fc, "H2II")
  762. my_fields.DI_density = get_field(fc, "DI")
  763. my_fields.DII_density = get_field(fc, "DII")
  764. my_fields.HDI_density = get_field(fc, "HDI")
  765. my_fields.e_density = get_field(fc, "de")
  766. my_fields.metal_density = get_field(fc, "metal")
  767. my_fields.dust_density = get_field(fc, "dust")
  768. my_fields.RT_heating_rate = get_field(fc, "RT_heating_rate")
  769. my_fields.volumetric_heating_rate = get_field(fc, "volumetric_heating_rate")
  770. my_fields.specific_heating_rate = get_field(fc, "specific_heating_rate")
  771. cdef gr_float *gamma = get_field(fc, "gamma")
  772.  
  773. c_local_calculate_gamma(
  774. chem_ptr,
  775. rates_ptr,
  776. units_ptr,
  777. &my_fields,
  778. gamma)
  779.  
  780. def calculate_pressure(fc):
  781. cdef c_chemistry_data *chem_ptr = NULL
  782. cdef c_chemistry_data_storage *rates_ptr = NULL
  783. cdef c_code_units *units_ptr = NULL
  784. _fetch_struct_ptrs(fc.chemistry_data, &chem_ptr, &rates_ptr, &units_ptr)
  785.  
  786. cdef int grid_dimension
  787. grid_dimension = fc["density"].shape[0]
  788. cdef np.ndarray ref_gs, ref_ge
  789. ref_gs = np.zeros(3, dtype="int32")
  790. ref_ge = np.zeros(3, dtype="int32")
  791. ref_ge[0] = grid_dimension -1
  792.  
  793. cdef c_field_data my_fields
  794. my_fields.grid_rank = 1
  795. my_fields.grid_dimension = &grid_dimension
  796. my_fields.grid_start = <int *> ref_gs.data
  797. my_fields.grid_end = <int *> ref_ge.data
  798. my_fields.density = get_field(fc, "density")
  799. my_fields.internal_energy = get_field(fc, "energy")
  800. my_fields.x_velocity = get_field(fc, "x-velocity")
  801. my_fields.y_velocity = get_field(fc, "y-velocity")
  802. my_fields.z_velocity = get_field(fc, "z-velocity")
  803. my_fields.HI_density = get_field(fc, "HI")
  804. my_fields.HII_density = get_field(fc, "HII")
  805. my_fields.HM_density = get_field(fc, "HM")
  806. my_fields.HeI_density = get_field(fc, "HeI")
  807. my_fields.HeII_density = get_field(fc, "HeII")
  808. my_fields.HeIII_density = get_field(fc, "HeIII")
  809. my_fields.H2I_density = get_field(fc, "H2I")
  810. my_fields.H2II_density = get_field(fc, "H2II")
  811. my_fields.DI_density = get_field(fc, "DI")
  812. my_fields.DII_density = get_field(fc, "DII")
  813. my_fields.HDI_density = get_field(fc, "HDI")
  814. my_fields.e_density = get_field(fc, "de")
  815. my_fields.metal_density = get_field(fc, "metal")
  816. my_fields.dust_density = get_field(fc, "dust")
  817. my_fields.RT_heating_rate = get_field(fc, "RT_heating_rate")
  818. my_fields.volumetric_heating_rate = get_field(fc, "volumetric_heating_rate")
  819. my_fields.specific_heating_rate = get_field(fc, "specific_heating_rate")
  820. cdef gr_float *pressure = get_field(fc, "pressure")
  821.  
  822. c_local_calculate_pressure(
  823. chem_ptr,
  824. rates_ptr,
  825. units_ptr,
  826. &my_fields,
  827. pressure)
  828.  
  829. def calculate_temperature(fc):
  830. cdef c_chemistry_data *chem_ptr = NULL
  831. cdef c_chemistry_data_storage *rates_ptr = NULL
  832. cdef c_code_units *units_ptr = NULL
  833. _fetch_struct_ptrs(fc.chemistry_data, &chem_ptr, &rates_ptr, &units_ptr)
  834.  
  835. cdef int grid_dimension
  836. grid_dimension = fc["density"].shape[0]
  837. cdef np.ndarray ref_gs, ref_ge
  838. ref_gs = np.zeros(3, dtype="int32")
  839. ref_ge = np.zeros(3, dtype="int32")
  840. ref_ge[0] = grid_dimension -1
  841.  
  842. cdef c_field_data my_fields
  843. my_fields.grid_rank = 1
  844. my_fields.grid_dimension = &grid_dimension
  845. my_fields.grid_start = <int *> ref_gs.data
  846. my_fields.grid_end = <int *> ref_ge.data
  847. my_fields.density = get_field(fc, "density")
  848. my_fields.internal_energy = get_field(fc, "energy")
  849. my_fields.x_velocity = get_field(fc, "x-velocity")
  850. my_fields.y_velocity = get_field(fc, "y-velocity")
  851. my_fields.z_velocity = get_field(fc, "z-velocity")
  852. my_fields.HI_density = get_field(fc, "HI")
  853. my_fields.HII_density = get_field(fc, "HII")
  854. my_fields.HM_density = get_field(fc, "HM")
  855. my_fields.HeI_density = get_field(fc, "HeI")
  856. my_fields.HeII_density = get_field(fc, "HeII")
  857. my_fields.HeIII_density = get_field(fc, "HeIII")
  858. my_fields.H2I_density = get_field(fc, "H2I")
  859. my_fields.H2II_density = get_field(fc, "H2II")
  860. my_fields.DI_density = get_field(fc, "DI")
  861. my_fields.DII_density = get_field(fc, "DII")
  862. my_fields.HDI_density = get_field(fc, "HDI")
  863. my_fields.e_density = get_field(fc, "de")
  864. my_fields.metal_density = get_field(fc, "metal")
  865. my_fields.dust_density = get_field(fc, "dust")
  866. my_fields.RT_heating_rate = get_field(fc, "RT_heating_rate")
  867. my_fields.volumetric_heating_rate = get_field(fc, "volumetric_heating_rate")
  868. my_fields.specific_heating_rate = get_field(fc, "specific_heating_rate")
  869. cdef gr_float *temperature = get_field(fc, "temperature")
  870.  
  871. c_local_calculate_temperature(
  872. chem_ptr,
  873. rates_ptr,
  874. units_ptr,
  875. &my_fields,
  876. temperature)
  877.  
  878. def calculate_dust_temperature(fc):
  879. cdef c_chemistry_data *chem_ptr = NULL
  880. cdef c_chemistry_data_storage *rates_ptr = NULL
  881. cdef c_code_units *units_ptr = NULL
  882. _fetch_struct_ptrs(fc.chemistry_data, &chem_ptr, &rates_ptr, &units_ptr)
  883.  
  884. cdef int grid_dimension
  885. grid_dimension = fc["density"].shape[0]
  886. cdef np.ndarray ref_gs, ref_ge
  887. ref_gs = np.zeros(3, dtype="int32")
  888. ref_ge = np.zeros(3, dtype="int32")
  889. ref_ge[0] = grid_dimension -1
  890.  
  891. cdef c_field_data my_fields
  892. my_fields.grid_rank = 1
  893. my_fields.grid_dimension = &grid_dimension
  894. my_fields.grid_start = <int *> ref_gs.data
  895. my_fields.grid_end = <int *> ref_ge.data
  896. my_fields.density = get_field(fc, "density")
  897. my_fields.internal_energy = get_field(fc, "energy")
  898. my_fields.x_velocity = get_field(fc, "x-velocity")
  899. my_fields.y_velocity = get_field(fc, "y-velocity")
  900. my_fields.z_velocity = get_field(fc, "z-velocity")
  901. my_fields.HI_density = get_field(fc, "HI")
  902. my_fields.HII_density = get_field(fc, "HII")
  903. my_fields.HM_density = get_field(fc, "HM")
  904. my_fields.HeI_density = get_field(fc, "HeI")
  905. my_fields.HeII_density = get_field(fc, "HeII")
  906. my_fields.HeIII_density = get_field(fc, "HeIII")
  907. my_fields.H2I_density = get_field(fc, "H2I")
  908. my_fields.H2II_density = get_field(fc, "H2II")
  909. my_fields.DI_density = get_field(fc, "DI")
  910. my_fields.DII_density = get_field(fc, "DII")
  911. my_fields.HDI_density = get_field(fc, "HDI")
  912. my_fields.e_density = get_field(fc, "de")
  913. my_fields.metal_density = get_field(fc, "metal")
  914. my_fields.dust_density = get_field(fc, "dust")
  915. my_fields.RT_heating_rate = get_field(fc, "RT_heating_rate")
  916. my_fields.volumetric_heating_rate = get_field(fc, "volumetric_heating_rate")
  917. my_fields.specific_heating_rate = get_field(fc, "specific_heating_rate")
  918. cdef gr_float *dust_temperature = get_field(fc, "dust_temperature")
  919.  
  920. c_local_calculate_dust_temperature(
  921. chem_ptr,
  922. rates_ptr,
  923. units_ptr,
  924. &my_fields,
  925. dust_temperature)
  926.  
  927. def get_grackle_version():
  928. cdef c_grackle_version version_struct = c_get_grackle_version()
  929. # all members of version_struct are string literals (i.e. don't call free)
  930. return {"version" : version_struct.version.decode('UTF-8'),
  931. "branch" : version_struct.branch.decode('UTF-8'),
  932. "revision" : version_struct.revision.decode('UTF-8')}
  933.  
  934. cdef list _get_parameter_name_list(const char*(*get_param)(size_t)):
  935. # the arg should be param_name_int, param_name_double, or param_name_string
  936.  
  937. cdef list out = []
  938. cdef const char* tmp
  939. cdef size_t i = 0
  940. while True:
  941. tmp = get_param(i)
  942. if tmp is NULL:
  943. return out
  944. else:
  945. out.append(tmp.decode('UTF-8'))
  946. i+=1
  947.  
  948. cdef class _wrapped_c_chemistry_data:
  949.  
  950. cdef c_chemistry_data data
  951. cdef dict _string_buffers
  952. # Each field in data that holds a string value, may have a corresponding
  953. # entry in _string_buffers (the key matches the name of the field).
  954. #
  955. # - if _string_buffers DOESN'T have an entry for <field>, then data.<field>
  956. # must currently store a default value. Depending on the specifics of
  957. # <field>, data.<field> holds either a NULL pointer or a pointer to a
  958. # string literal (stored in the C library).
  959. #
  960. # - if _string_buffers DOES have an entry for <field>, then data.<field>
  961. # points to the buffer of the Python bytes object given by
  962. # _string_buffers["<field>"]. Note that the string is always terminated
  963. # by character (the Python interface just hides if from users)
  964.  
  965. def __cinit__(self):
  966. self.data = _set_default_chemistry_parameters()
  967. self._string_buffers = {}
  968.  
  969. def _access_struct_field(self, key, val = None):
  970. """
  971. Return a copy of the value of the field stored in the wrapped struct.
  972.  
  973. When val is not None, this function updates the stored value first (the
  974. returned value is effectively a copy of val after some type coercion)
  975. """
  976. cdef c_chemistry_data* data_ptr = &self.data
  977.  
  978. # determine whether the field needs to be updated
  979. update_field = (val is not None)
  980.  
  981. # coerce key to a bytes object and then get the analogous null
  982. # terminated string to pass to the C functions
  983. coerced_key = key.encode('ascii') if isinstance(key, str) else key
  984. cdef char* _key = coerced_key # _key's lifetime is tied to coerced_key
  985.  
  986. # handle the case where key is an integer
  987. cdef int* int_p = local_chemistry_data_access_int(data_ptr, _key)
  988. if int_p is not NULL:
  989. if update_field:
  990. coerced_val = int(<int?>val)
  991. if coerced_val != val: # prevent invalid assignment of a
  992. # non-integer floating point
  993. raise TypeError("must be convertable to an int without a "
  994. "loss of information.")
  995. int_p[0] = <int>coerced_val
  996. return int(int_p[0])
  997.  
  998. # handle the case where key is a double
  999. cdef double* dbl_p = local_chemistry_data_access_double(data_ptr, _key)
  1000. if dbl_p is not NULL:
  1001. if update_field:
  1002. dbl_p[0] = <double?>val # <double?> cast performs type checking
  1003. return float(dbl_p[0])
  1004.  
  1005. # handle the case where key is a string
  1006. cdef char* tmp = NULL
  1007. cdef char** str_p = local_chemistry_data_access_string(data_ptr, _key)
  1008. if (str_p is not NULL): # this case requires additional care
  1009.  
  1010. if update_field:
  1011. # coerce val to a bytes object & store in self._string_buffers
  1012. if isinstance(val, str):
  1013. self._string_buffers[key] = bytes(val.encode('ascii'))
  1014. elif isinstance(val, (bytes,bytearray)):
  1015. self._string_buffers[key] = bytes(val)
  1016. else:
  1017. raise TypeError(type(val))
  1018.  
  1019. assert b'\0' not in self._string_buffers[key] # sanity check
  1020. tmp = <bytes>self._string_buffers[key]
  1021. str_p[0] = tmp
  1022.  
  1023. elif (str_p[0] is NULL):
  1024. # handle the case where we are accessing a default field value
  1025. # and it was initialized to a NULL pointer
  1026. return None
  1027.  
  1028. # for simplicity, just return a new copy of the contained string.
  1029. # (this correctly handles the case where str_p[0] holds a pointer
  1030. # to a string literal managed by the c library)
  1031. return bytes(str_p[0])
  1032.  
  1033. raise KeyError(key) # the key is not recognized
  1034.  
  1035. def __getitem__(self, key):
  1036. return self._access_struct_field(key, val = None)
  1037.  
  1038. def __setitem__(self, key, value):
  1039. if value is None:
  1040. raise ValueError("Members of c_chemistry_data are not allowed to "
  1041. "be assigned values of None")
  1042. self._access_struct_field(key, val = value)
  1043.  
  1044. @staticmethod
  1045. def int_keys(): return _get_parameter_name_list(param_name_int)
  1046.  
  1047. @staticmethod
  1048. def double_keys(): return _get_parameter_name_list(param_name_double)
  1049.  
  1050. @staticmethod
  1051. def string_keys(): return _get_parameter_name_list(param_name_string)
  1052.  
  1053. def __iter__(self):
  1054. return iter(self.int_keys() + self.double_keys() + self.string_keys())
  1055.  
  1056. def __len__(self):
  1057. return (len(self.int_keys()) + len(self.double_keys()) +
  1058. len(self.string_keys()))
  1059.  
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement
Public Pastes
Advertisement
We use cookies for various purposes including analytics. By continuing to use Pastebin, you agree to our use of cookies as described in the Cookies Policy.  OK, I Understand
Not a member of Pastebin yet?
Sign Up, it unlocks many cool features!