roms-tools 1.6.2__py3-none-any.whl → 1.7.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 (221) hide show
  1. ci/environment.yml +1 -1
  2. roms_tools/__init__.py +1 -0
  3. roms_tools/_version.py +1 -1
  4. roms_tools/setup/boundary_forcing.py +13 -112
  5. roms_tools/setup/datasets.py +778 -191
  6. roms_tools/setup/download.py +30 -0
  7. roms_tools/setup/initial_conditions.py +14 -76
  8. roms_tools/setup/plot.py +77 -15
  9. roms_tools/setup/river_forcing.py +589 -0
  10. roms_tools/setup/surface_forcing.py +10 -112
  11. roms_tools/setup/tides.py +6 -67
  12. roms_tools/setup/utils.py +259 -1
  13. roms_tools/tests/test_setup/test_boundary_forcing.py +0 -2
  14. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/.zattrs +1 -1
  15. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/.zmetadata +157 -130
  16. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/ALK_ALT_CO2_east/.zattrs +1 -1
  17. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/ALK_ALT_CO2_north/.zattrs +1 -1
  18. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/ALK_ALT_CO2_south/.zattrs +1 -1
  19. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/ALK_ALT_CO2_west/.zattrs +1 -1
  20. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/ALK_east/.zattrs +1 -1
  21. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/ALK_north/.zattrs +1 -1
  22. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/ALK_south/.zattrs +1 -1
  23. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/ALK_west/.zattrs +1 -1
  24. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/DIC_ALT_CO2_east/.zattrs +1 -1
  25. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/DIC_ALT_CO2_north/.zattrs +1 -1
  26. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/DIC_ALT_CO2_south/.zattrs +1 -1
  27. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/DIC_ALT_CO2_west/.zattrs +1 -1
  28. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/DIC_east/.zattrs +1 -1
  29. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/DIC_north/.zattrs +1 -1
  30. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/DIC_south/.zattrs +1 -1
  31. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/DIC_west/.zattrs +1 -1
  32. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/DOC_east/.zattrs +1 -1
  33. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/DOC_north/.zattrs +1 -1
  34. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/DOC_south/.zattrs +1 -1
  35. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/DOC_west/.zattrs +1 -1
  36. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/DOCr_east/.zattrs +1 -1
  37. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/DOCr_north/.zattrs +1 -1
  38. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/DOCr_south/.zattrs +1 -1
  39. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/DOCr_west/.zattrs +1 -1
  40. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/DON_east/.zattrs +1 -1
  41. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/DON_north/.zattrs +1 -1
  42. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/DON_south/.zattrs +1 -1
  43. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/DON_west/.zattrs +1 -1
  44. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/DONr_east/.zattrs +1 -1
  45. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/DONr_north/.zattrs +1 -1
  46. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/DONr_south/.zattrs +1 -1
  47. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/DONr_west/.zattrs +1 -1
  48. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/DOP_east/.zattrs +1 -1
  49. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/DOP_north/.zattrs +1 -1
  50. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/DOP_south/.zattrs +1 -1
  51. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/DOP_west/.zattrs +1 -1
  52. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/DOPr_east/.zattrs +1 -1
  53. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/DOPr_north/.zattrs +1 -1
  54. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/DOPr_south/.zattrs +1 -1
  55. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/DOPr_west/.zattrs +1 -1
  56. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/Fe_east/.zattrs +1 -1
  57. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/Fe_north/.zattrs +1 -1
  58. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/Fe_south/.zattrs +1 -1
  59. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/Fe_west/.zattrs +1 -1
  60. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/Lig_east/.zattrs +1 -1
  61. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/Lig_north/.zattrs +1 -1
  62. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/Lig_south/.zattrs +1 -1
  63. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/Lig_west/.zattrs +1 -1
  64. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/NH4_east/.zattrs +1 -1
  65. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/NH4_north/.zattrs +1 -1
  66. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/NH4_south/.zattrs +1 -1
  67. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/NH4_west/.zattrs +1 -1
  68. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/NO3_east/.zattrs +1 -1
  69. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/NO3_north/.zattrs +1 -1
  70. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/NO3_south/.zattrs +1 -1
  71. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/NO3_west/.zattrs +1 -1
  72. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/O2_east/.zattrs +1 -1
  73. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/O2_north/.zattrs +1 -1
  74. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/O2_south/.zattrs +1 -1
  75. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/O2_west/.zattrs +1 -1
  76. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/PO4_east/.zattrs +1 -1
  77. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/PO4_north/.zattrs +1 -1
  78. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/PO4_south/.zattrs +1 -1
  79. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/PO4_west/.zattrs +1 -1
  80. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/SiO3_east/.zattrs +1 -1
  81. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/SiO3_north/.zattrs +1 -1
  82. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/SiO3_south/.zattrs +1 -1
  83. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/SiO3_west/.zattrs +1 -1
  84. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/abs_time/.zattrs +1 -0
  85. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/bry_time/.zattrs +1 -1
  86. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/diatC_east/.zattrs +1 -1
  87. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/diatC_north/.zattrs +1 -1
  88. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/diatC_south/.zattrs +1 -1
  89. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/diatC_west/.zattrs +1 -1
  90. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/diatChl_east/.zattrs +1 -1
  91. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/diatChl_north/.zattrs +1 -1
  92. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/diatChl_south/.zattrs +1 -1
  93. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/diatChl_west/.zattrs +1 -1
  94. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/diatFe_east/.zattrs +1 -1
  95. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/diatFe_north/.zattrs +1 -1
  96. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/diatFe_south/.zattrs +1 -1
  97. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/diatFe_west/.zattrs +1 -1
  98. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/diatP_east/.zattrs +1 -1
  99. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/diatP_north/.zattrs +1 -1
  100. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/diatP_south/.zattrs +1 -1
  101. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/diatP_west/.zattrs +1 -1
  102. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/diatSi_east/.zattrs +1 -1
  103. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/diatSi_north/.zattrs +1 -1
  104. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/diatSi_south/.zattrs +1 -1
  105. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/diatSi_west/.zattrs +1 -1
  106. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/diazC_east/.zattrs +1 -1
  107. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/diazC_north/.zattrs +1 -1
  108. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/diazC_south/.zattrs +1 -1
  109. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/diazC_west/.zattrs +1 -1
  110. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/diazChl_east/.zattrs +1 -1
  111. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/diazChl_north/.zattrs +1 -1
  112. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/diazChl_south/.zattrs +1 -1
  113. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/diazChl_west/.zattrs +1 -1
  114. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/diazFe_east/.zattrs +1 -1
  115. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/diazFe_north/.zattrs +1 -1
  116. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/diazFe_south/.zattrs +1 -1
  117. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/diazFe_west/.zattrs +1 -1
  118. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/diazP_east/.zattrs +1 -1
  119. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/diazP_north/.zattrs +1 -1
  120. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/diazP_south/.zattrs +1 -1
  121. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/diazP_west/.zattrs +1 -1
  122. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/month/.zarray +20 -0
  123. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/month/.zattrs +6 -0
  124. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/month/0 +0 -0
  125. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/spC_east/.zattrs +1 -1
  126. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/spC_north/.zattrs +1 -1
  127. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/spC_south/.zattrs +1 -1
  128. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/spC_west/.zattrs +1 -1
  129. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/spCaCO3_east/.zattrs +1 -1
  130. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/spCaCO3_north/.zattrs +1 -1
  131. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/spCaCO3_south/.zattrs +1 -1
  132. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/spCaCO3_west/.zattrs +1 -1
  133. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/spChl_east/.zattrs +1 -1
  134. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/spChl_north/.zattrs +1 -1
  135. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/spChl_south/.zattrs +1 -1
  136. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/spChl_west/.zattrs +1 -1
  137. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/spFe_east/.zattrs +1 -1
  138. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/spFe_north/.zattrs +1 -1
  139. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/spFe_south/.zattrs +1 -1
  140. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/spFe_west/.zattrs +1 -1
  141. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/spP_east/.zattrs +1 -1
  142. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/spP_north/.zattrs +1 -1
  143. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/spP_south/.zattrs +1 -1
  144. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/spP_west/.zattrs +1 -1
  145. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/zooC_east/.zattrs +1 -1
  146. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/zooC_north/.zattrs +1 -1
  147. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/zooC_south/.zattrs +1 -1
  148. roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/zooC_west/.zattrs +1 -1
  149. roms_tools/tests/test_setup/test_data/bgc_surface_forcing_from_climatology.zarr/.zattrs +1 -1
  150. roms_tools/tests/test_setup/test_data/bgc_surface_forcing_from_climatology.zarr/.zmetadata +39 -12
  151. roms_tools/tests/test_setup/test_data/bgc_surface_forcing_from_climatology.zarr/abs_time/.zattrs +1 -0
  152. roms_tools/tests/test_setup/test_data/bgc_surface_forcing_from_climatology.zarr/dust/.zattrs +1 -1
  153. roms_tools/tests/test_setup/test_data/bgc_surface_forcing_from_climatology.zarr/dust_time/.zattrs +1 -1
  154. roms_tools/tests/test_setup/test_data/bgc_surface_forcing_from_climatology.zarr/iron/.zattrs +1 -1
  155. roms_tools/tests/test_setup/test_data/bgc_surface_forcing_from_climatology.zarr/iron_time/.zattrs +1 -1
  156. roms_tools/tests/test_setup/test_data/bgc_surface_forcing_from_climatology.zarr/month/.zarray +20 -0
  157. roms_tools/tests/test_setup/test_data/bgc_surface_forcing_from_climatology.zarr/month/.zattrs +6 -0
  158. roms_tools/tests/test_setup/test_data/bgc_surface_forcing_from_climatology.zarr/month/0 +0 -0
  159. roms_tools/tests/test_setup/test_data/bgc_surface_forcing_from_climatology.zarr/nhy/.zattrs +1 -1
  160. roms_tools/tests/test_setup/test_data/bgc_surface_forcing_from_climatology.zarr/nhy_time/.zattrs +1 -1
  161. roms_tools/tests/test_setup/test_data/bgc_surface_forcing_from_climatology.zarr/nox/.zattrs +1 -1
  162. roms_tools/tests/test_setup/test_data/bgc_surface_forcing_from_climatology.zarr/nox_time/.zattrs +1 -1
  163. roms_tools/tests/test_setup/test_data/bgc_surface_forcing_from_climatology.zarr/pco2_air/.zattrs +1 -1
  164. roms_tools/tests/test_setup/test_data/bgc_surface_forcing_from_climatology.zarr/pco2_air_alt/.zattrs +1 -1
  165. roms_tools/tests/test_setup/test_data/bgc_surface_forcing_from_climatology.zarr/pco2_time/.zattrs +1 -1
  166. roms_tools/tests/test_setup/test_data/river_forcing.zarr/.zattrs +3 -0
  167. roms_tools/tests/test_setup/test_data/river_forcing.zarr/.zgroup +3 -0
  168. roms_tools/tests/test_setup/test_data/river_forcing.zarr/.zmetadata +214 -0
  169. roms_tools/tests/test_setup/test_data/river_forcing.zarr/abs_time/.zarray +20 -0
  170. roms_tools/tests/test_setup/test_data/river_forcing.zarr/abs_time/.zattrs +8 -0
  171. roms_tools/tests/test_setup/test_data/river_forcing.zarr/abs_time/0 +0 -0
  172. roms_tools/tests/test_setup/test_data/river_forcing.zarr/month/.zarray +20 -0
  173. roms_tools/tests/test_setup/test_data/river_forcing.zarr/month/.zattrs +6 -0
  174. roms_tools/tests/test_setup/test_data/river_forcing.zarr/month/0 +0 -0
  175. roms_tools/tests/test_setup/test_data/river_forcing.zarr/river_name/.zarray +24 -0
  176. roms_tools/tests/test_setup/test_data/river_forcing.zarr/river_name/.zattrs +6 -0
  177. roms_tools/tests/test_setup/test_data/river_forcing.zarr/river_name/0 +0 -0
  178. roms_tools/tests/test_setup/test_data/river_forcing.zarr/river_time/.zarray +20 -0
  179. roms_tools/tests/test_setup/test_data/river_forcing.zarr/river_time/.zattrs +8 -0
  180. roms_tools/tests/test_setup/test_data/river_forcing.zarr/river_time/0 +0 -0
  181. roms_tools/tests/test_setup/test_data/river_forcing.zarr/river_tracer/.zarray +24 -0
  182. roms_tools/tests/test_setup/test_data/river_forcing.zarr/river_tracer/.zattrs +10 -0
  183. roms_tools/tests/test_setup/test_data/river_forcing.zarr/river_tracer/0.0.0 +0 -0
  184. roms_tools/tests/test_setup/test_data/river_forcing.zarr/river_volume/.zarray +22 -0
  185. roms_tools/tests/test_setup/test_data/river_forcing.zarr/river_volume/.zattrs +9 -0
  186. roms_tools/tests/test_setup/test_data/river_forcing.zarr/river_volume/0.0 +0 -0
  187. roms_tools/tests/test_setup/test_data/river_forcing.zarr/tracer_name/.zarray +20 -0
  188. roms_tools/tests/test_setup/test_data/river_forcing.zarr/tracer_name/.zattrs +6 -0
  189. roms_tools/tests/test_setup/test_data/river_forcing.zarr/tracer_name/0 +0 -0
  190. roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/.zattrs +1 -0
  191. roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/.zgroup +3 -0
  192. roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/.zmetadata +185 -0
  193. roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/abs_time/.zarray +20 -0
  194. roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/abs_time/.zattrs +8 -0
  195. roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/abs_time/0 +0 -0
  196. roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/river_name/.zarray +24 -0
  197. roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/river_name/.zattrs +6 -0
  198. roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/river_name/0 +0 -0
  199. roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/river_time/.zarray +20 -0
  200. roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/river_time/.zattrs +7 -0
  201. roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/river_time/0 +0 -0
  202. roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/river_tracer/.zarray +24 -0
  203. roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/river_tracer/.zattrs +10 -0
  204. roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/river_tracer/0.0.0 +0 -0
  205. roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/river_volume/.zarray +22 -0
  206. roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/river_volume/.zattrs +9 -0
  207. roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/river_volume/0.0 +0 -0
  208. roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/tracer_name/.zarray +20 -0
  209. roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/tracer_name/.zattrs +6 -0
  210. roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/tracer_name/0 +0 -0
  211. roms_tools/tests/test_setup/test_initial_conditions.py +0 -2
  212. roms_tools/tests/test_setup/test_river_forcing.py +366 -0
  213. roms_tools/tests/test_setup/test_surface_forcing.py +0 -2
  214. roms_tools/tests/test_setup/test_tides.py +0 -2
  215. roms_tools/tests/test_setup/test_validation.py +4 -0
  216. roms_tools/utils.py +12 -10
  217. {roms_tools-1.6.2.dist-info → roms_tools-1.7.0.dist-info}/METADATA +5 -5
  218. {roms_tools-1.6.2.dist-info → roms_tools-1.7.0.dist-info}/RECORD +221 -168
  219. {roms_tools-1.6.2.dist-info → roms_tools-1.7.0.dist-info}/WHEEL +1 -1
  220. {roms_tools-1.6.2.dist-info → roms_tools-1.7.0.dist-info}/LICENSE +0 -0
  221. {roms_tools-1.6.2.dist-info → roms_tools-1.7.0.dist-info}/top_level.txt +0 -0
@@ -13,10 +13,13 @@ from roms_tools.setup.utils import (
13
13
  get_time_type,
14
14
  convert_cftime_to_datetime,
15
15
  one_dim_fill,
16
+ gc_dist,
16
17
  )
17
- from roms_tools.setup.download import download_correction_data
18
+ from roms_tools.setup.download import download_correction_data, download_river_data
18
19
  from roms_tools.setup.fill import LateralFill
19
20
 
21
+ # lat-lon datasets
22
+
20
23
 
21
24
  @dataclass(frozen=True, kw_only=True)
22
25
  class Dataset:
@@ -32,10 +35,10 @@ class Dataset:
32
35
  end_time : Optional[datetime], optional
33
36
  The end time for selecting relevant data. If not provided, only data at the start_time is selected if start_time is provided,
34
37
  or no filtering is applied if start_time is not provided.
35
- var_names: Dict[str, str]
36
- Dictionary of variable names that are required in the dataset.
37
38
  dim_names: Dict[str, str], optional
38
39
  Dictionary specifying the names of dimensions in the dataset.
40
+ var_names: Dict[str, str]
41
+ Dictionary of variable names that are required in the dataset.
39
42
  climatology : bool
40
43
  Indicates whether the dataset is climatological. Defaults to False.
41
44
  use_dask: bool
@@ -62,7 +65,6 @@ class Dataset:
62
65
  filename: Union[str, Path, List[Union[str, Path]]]
63
66
  start_time: Optional[datetime] = None
64
67
  end_time: Optional[datetime] = None
65
- var_names: Dict[str, str]
66
68
  dim_names: Dict[str, str] = field(
67
69
  default_factory=lambda: {
68
70
  "longitude": "longitude",
@@ -70,6 +72,7 @@ class Dataset:
70
72
  "time": "time",
71
73
  }
72
74
  )
75
+ var_names: Dict[str, str]
73
76
  climatology: Optional[bool] = False
74
77
  use_dask: Optional[bool] = True
75
78
  apply_post_processing: Optional[bool] = True
@@ -149,101 +152,7 @@ class Dataset:
149
152
  If a list of files is provided but self.dim_names["time"] is not available or use_dask=False.
150
153
  """
151
154
 
152
- # Precompile the regex for matching wildcard characters
153
- wildcard_regex = re.compile(r"[\*\?\[\]]")
154
-
155
- # Convert Path objects to strings
156
- if isinstance(self.filename, (str, Path)):
157
- filename_str = str(self.filename)
158
- elif isinstance(self.filename, list):
159
- filename_str = [str(f) for f in self.filename]
160
- else:
161
- raise ValueError(
162
- "filename must be a string, Path, or a list of strings/Paths."
163
- )
164
-
165
- # Handle the case when filename is a string
166
- contains_wildcard = False
167
- if isinstance(filename_str, str):
168
- contains_wildcard = bool(wildcard_regex.search(filename_str))
169
- if contains_wildcard:
170
- matching_files = glob.glob(filename_str)
171
- if not matching_files:
172
- raise FileNotFoundError(
173
- f"No files found matching the pattern '{filename_str}'."
174
- )
175
- else:
176
- matching_files = [filename_str]
177
-
178
- # Handle the case when filename is a list
179
- elif isinstance(filename_str, list):
180
- contains_wildcard = any(wildcard_regex.search(f) for f in filename_str)
181
- if contains_wildcard:
182
- matching_files = []
183
- for f in filename_str:
184
- files = glob.glob(f)
185
- if not files:
186
- raise FileNotFoundError(
187
- f"No files found matching the pattern '{f}'."
188
- )
189
- matching_files.extend(files)
190
- else:
191
- matching_files = filename_str
192
-
193
- # Check if time dimension is available when multiple files are provided
194
- if isinstance(filename_str, list) and "time" not in self.dim_names:
195
- raise ValueError(
196
- "A list of files is provided, but time dimension is not available. "
197
- "A time dimension must be available to concatenate the files."
198
- )
199
-
200
- # Determine the kwargs for combining datasets
201
- if contains_wildcard or len(matching_files) == 1:
202
- # If there is a wildcard or just one file, use by_coords
203
- kwargs = {"combine": "by_coords"}
204
- else:
205
- # Otherwise, use nested combine based on time
206
- kwargs = {"combine": "nested", "concat_dim": self.dim_names["time"]}
207
-
208
- # Base kwargs used for dataset combination
209
- combine_kwargs = {
210
- "coords": "minimal",
211
- "compat": "override",
212
- "combine_attrs": "override",
213
- }
214
-
215
- if self.use_dask:
216
-
217
- chunks = {
218
- self.dim_names["latitude"]: -1,
219
- self.dim_names["longitude"]: -1,
220
- }
221
- if "depth" in self.dim_names:
222
- chunks[self.dim_names["depth"]] = -1
223
- if "time" in self.dim_names:
224
- chunks[self.dim_names["time"]] = 1
225
-
226
- ds = xr.open_mfdataset(
227
- matching_files,
228
- chunks=chunks,
229
- **combine_kwargs,
230
- **kwargs,
231
- )
232
- else:
233
- ds_list = []
234
- for file in matching_files:
235
- ds = xr.open_dataset(file, chunks=None)
236
- ds_list.append(ds)
237
-
238
- if kwargs["combine"] == "by_coords":
239
- ds = xr.combine_by_coords(ds_list, **combine_kwargs)
240
- elif kwargs["combine"] == "nested":
241
- ds = xr.combine_nested(
242
- ds_list, concat_dim=kwargs["concat_dim"], **combine_kwargs
243
- )
244
-
245
- if "time" in self.dim_names and self.dim_names["time"] not in ds.dims:
246
- ds = ds.expand_dims(self.dim_names["time"])
155
+ ds = _load_data(self.filename, self.dim_names, self.use_dask)
247
156
 
248
157
  return ds
249
158
 
@@ -278,19 +187,8 @@ class Dataset:
278
187
  ValueError
279
188
  If the dataset does not contain the specified variables or dimensions.
280
189
  """
281
- missing_vars = [
282
- var for var in self.var_names.values() if var not in ds.data_vars
283
- ]
284
- if missing_vars:
285
- raise ValueError(
286
- f"Dataset does not contain all required variables. The following variables are missing: {missing_vars}"
287
- )
288
190
 
289
- missing_dims = [dim for dim in self.dim_names.values() if dim not in ds.dims]
290
- if missing_dims:
291
- raise ValueError(
292
- f"Dataset does not contain all required dimensions. The following dimensions are missing: {missing_vars}"
293
- )
191
+ _check_dataset(ds, self.dim_names, self.var_names)
294
192
 
295
193
  def select_relevant_fields(self, ds) -> xr.Dataset:
296
194
  """Selects and returns a subset of the dataset containing only the variables
@@ -379,86 +277,10 @@ class Dataset:
379
277
  """
380
278
 
381
279
  time_dim = self.dim_names["time"]
382
- if time_dim in ds.variables:
383
- if self.climatology:
384
- if len(ds[time_dim]) != 12:
385
- raise ValueError(
386
- f"The dataset contains {len(ds[time_dim])} time steps, but the climatology flag is set to True, which requires exactly 12 time steps."
387
- )
388
- if not self.end_time:
389
- # Interpolate from climatology for initial conditions
390
- ds = interpolate_from_climatology(
391
- ds, self.dim_names["time"], self.start_time
392
- )
393
- else:
394
- time_type = get_time_type(ds[time_dim])
395
- if time_type == "int":
396
- raise ValueError(
397
- "The dataset contains integer time values, which are only supported when the climatology flag is set to True. However, your climatology flag is set to False."
398
- )
399
- if time_type == "cftime":
400
- ds = ds.assign_coords(
401
- {time_dim: convert_cftime_to_datetime(ds[time_dim])}
402
- )
403
- if self.end_time:
404
- end_time = self.end_time
405
-
406
- # Identify records before or at start_time
407
- before_start = ds[time_dim] <= np.datetime64(self.start_time)
408
- if before_start.any():
409
- closest_before_start = (
410
- ds[time_dim].where(before_start, drop=True).max()
411
- )
412
- else:
413
- logging.warning("No records found at or before the start_time.")
414
- closest_before_start = ds[time_dim].min()
415
-
416
- # Identify records after or at end_time
417
- after_end = ds[time_dim] >= np.datetime64(end_time)
418
- if after_end.any():
419
- closest_after_end = (
420
- ds[time_dim].where(after_end, drop=True).min()
421
- )
422
- else:
423
- logging.warning("No records found at or after the end_time.")
424
- closest_after_end = ds[time_dim].max()
425
-
426
- # Select records within the time range and add the closest before/after
427
- within_range = (ds[time_dim] > np.datetime64(self.start_time)) & (
428
- ds[time_dim] < np.datetime64(end_time)
429
- )
430
- selected_times = ds[time_dim].where(
431
- within_range
432
- | (ds[time_dim] == closest_before_start)
433
- | (ds[time_dim] == closest_after_end),
434
- drop=True,
435
- )
436
- ds = ds.sel({time_dim: selected_times})
437
- else:
438
- # Look in time range [self.start_time, self.start_time + 24h]
439
- end_time = self.start_time + timedelta(days=1)
440
- times = (np.datetime64(self.start_time) <= ds[time_dim]) & (
441
- ds[time_dim] < np.datetime64(end_time)
442
- )
443
- if np.all(~times):
444
- raise ValueError(
445
- f"The dataset does not contain any time entries between the specified start_time: {self.start_time} "
446
- f"and {self.start_time + timedelta(hours=24)}. "
447
- "Please ensure the dataset includes time entries for that range."
448
- )
449
-
450
- ds = ds.where(times, drop=True)
451
- if ds.sizes[time_dim] > 1:
452
- # Pick the time closest to self.start_time
453
- ds = ds.isel({time_dim: 0})
454
- logging.info(
455
- f"Selected time entry closest to the specified start_time ({self.start_time}) within the range [{self.start_time}, {self.start_time + timedelta(hours=24)}]: {ds[time_dim].values}"
456
- )
457
- else:
458
- logging.warning(
459
- "Dataset does not contain any time information. Please check if the time dimension "
460
- "is correctly named or if the dataset includes time data."
461
- )
280
+
281
+ ds = _select_relevant_times(
282
+ ds, time_dim, self.start_time, self.end_time, self.climatology
283
+ )
462
284
 
463
285
  return ds
464
286
 
@@ -1522,3 +1344,768 @@ class ERA5Correction(Dataset):
1522
1344
  "The correction dataset does not contain all specified longitude values."
1523
1345
  )
1524
1346
  object.__setattr__(self, "ds", subdomain)
1347
+
1348
+
1349
+ # river datasets
1350
+ @dataclass(frozen=True, kw_only=True)
1351
+ class RiverDataset:
1352
+ """Represents river data.
1353
+
1354
+ Parameters
1355
+ ----------
1356
+ filename : Union[str, Path, List[Union[str, Path]]]
1357
+ The path to the data file(s). Can be a single string (with or without wildcards), a single Path object,
1358
+ or a list of strings or Path objects containing multiple files.
1359
+ start_time : datetime
1360
+ The start time for selecting relevant data.
1361
+ end_time : datetime
1362
+ The end time for selecting relevant data.
1363
+ dim_names: Dict[str, str]
1364
+ Dictionary specifying the names of dimensions in the dataset.
1365
+ Requires "station" and "time" as keys.
1366
+ var_names: Dict[str, str]
1367
+ Dictionary of variable names that are required in the dataset.
1368
+ Requires the keys "latitude", "longitude", "flux", "ratio", and "name".
1369
+ opt_var_names: Dict[str, str], optional
1370
+ Dictionary of variable names that are optional in the dataset.
1371
+ Defaults to an empty dictionary.
1372
+ climatology : bool
1373
+ Indicates whether the dataset is climatological. Defaults to False.
1374
+
1375
+ Attributes
1376
+ ----------
1377
+ ds : xr.Dataset
1378
+ The xarray Dataset containing the forcing data on its original grid.
1379
+ """
1380
+
1381
+ filename: Union[str, Path, List[Union[str, Path]]]
1382
+ start_time: datetime
1383
+ end_time: datetime
1384
+ dim_names: Dict[str, str]
1385
+ var_names: Dict[str, str]
1386
+ opt_var_names: Optional[Dict[str, str]] = field(default_factory=dict)
1387
+ climatology: Optional[bool] = False
1388
+ ds: xr.Dataset = field(init=False, repr=False)
1389
+
1390
+ def __post_init__(self):
1391
+
1392
+ # Validate start_time and end_time
1393
+ if not isinstance(self.start_time, datetime):
1394
+ raise TypeError(
1395
+ f"start_time must be a datetime object, but got {type(self.start_time).__name__}."
1396
+ )
1397
+ if not isinstance(self.end_time, datetime):
1398
+ raise TypeError(
1399
+ f"end_time must be a datetime object, but got {type(self.end_time).__name__}."
1400
+ )
1401
+
1402
+ ds = self.load_data()
1403
+ ds = self.clean_up(ds)
1404
+ self.check_dataset(ds)
1405
+
1406
+ # Select relevant times
1407
+ ds = self.add_time_info(ds)
1408
+ object.__setattr__(self, "ds", ds)
1409
+
1410
+ def load_data(self) -> xr.Dataset:
1411
+ """Load dataset from the specified file.
1412
+
1413
+ Returns
1414
+ -------
1415
+ ds : xr.Dataset
1416
+ The loaded xarray Dataset containing the forcing data.
1417
+
1418
+ Raises
1419
+ ------
1420
+ FileNotFoundError
1421
+ If the specified file does not exist.
1422
+ ValueError
1423
+ If a list of files is provided but self.dim_names["time"] is not available or use_dask=False.
1424
+ """
1425
+ ds = _load_data(
1426
+ self.filename, self.dim_names, use_dask=False, decode_times=False
1427
+ )
1428
+
1429
+ return ds
1430
+
1431
+ def clean_up(self, ds: xr.Dataset) -> xr.Dataset:
1432
+ """Decodes the 'name' variable (if byte-encoded) and updates the dataset.
1433
+
1434
+ This method checks if the 'name' variable is of dtype 'object' (i.e., byte-encoded),
1435
+ and if so, decodes each byte array to a string and updates the dataset.
1436
+ It also ensures that the 'station' dimension is of integer type.
1437
+
1438
+
1439
+ Parameters
1440
+ ----------
1441
+ ds : xr.Dataset
1442
+ The dataset containing the 'name' variable to decode.
1443
+
1444
+ Returns
1445
+ -------
1446
+ ds : xr.Dataset
1447
+ The dataset with the decoded 'name' variable.
1448
+ """
1449
+
1450
+ if ds[self.var_names["name"]].dtype == "object":
1451
+ names = []
1452
+ for i in range(len(ds[self.dim_names["station"]])):
1453
+ byte_array = ds[self.var_names["name"]].isel(
1454
+ **{self.dim_names["station"]: i}
1455
+ )
1456
+ name = decode_string(byte_array)
1457
+ names.append(name)
1458
+ ds[self.var_names["name"]] = xr.DataArray(
1459
+ data=names, dims=self.dim_names["station"]
1460
+ )
1461
+
1462
+ if ds[self.dim_names["station"]].dtype == "float64":
1463
+ ds[self.dim_names["station"]] = ds[self.dim_names["station"]].astype(int)
1464
+
1465
+ # Drop all variables that have chars dim
1466
+ vars_to_drop = ["ocn_name", "stn_name", "ct_name", "cn_name", "chars"]
1467
+ existing_vars = [var for var in vars_to_drop if var in ds]
1468
+ ds = ds.drop_vars(existing_vars)
1469
+
1470
+ return ds
1471
+
1472
+ def check_dataset(self, ds: xr.Dataset) -> None:
1473
+ """Check if the dataset contains the specified variables and dimensions.
1474
+
1475
+ Parameters
1476
+ ----------
1477
+ ds : xr.Dataset
1478
+ The xarray Dataset to check.
1479
+
1480
+ Raises
1481
+ ------
1482
+ ValueError
1483
+ If the dataset does not contain the specified variables or dimensions.
1484
+ """
1485
+
1486
+ _check_dataset(ds, self.dim_names, self.var_names, self.opt_var_names)
1487
+
1488
+ def add_time_info(self, ds: xr.Dataset) -> xr.Dataset:
1489
+ """Dummy method to be overridden by child classes to add time information to the
1490
+ dataset.
1491
+
1492
+ This method is intended as a placeholder and should be implemented in subclasses
1493
+ to provide specific functionality for adding time-related information to the dataset.
1494
+
1495
+ Parameters
1496
+ ----------
1497
+ ds : xr.Dataset
1498
+ The xarray Dataset to which time information will be added.
1499
+
1500
+ Returns
1501
+ -------
1502
+ xr.Dataset
1503
+ The xarray Dataset with time information added (as implemented by child classes).
1504
+ """
1505
+ return ds
1506
+
1507
+ def select_relevant_times(self, ds) -> xr.Dataset:
1508
+ """Select a subset of the dataset based on the specified time range.
1509
+
1510
+ This method filters the dataset to include all records between `start_time` and `end_time`.
1511
+ Additionally, it ensures that one record at or before `start_time` and one record at or
1512
+ after `end_time` are included, even if they fall outside the strict time range.
1513
+
1514
+ If no `end_time` is specified, the method will select the time range of
1515
+ [start_time, start_time + 24 hours] and return the closest time entry to `start_time` within that range.
1516
+
1517
+ Parameters
1518
+ ----------
1519
+ ds : xr.Dataset
1520
+ The input dataset to be filtered. Must contain a time dimension.
1521
+
1522
+ Returns
1523
+ -------
1524
+ xr.Dataset
1525
+ A dataset filtered to the specified time range, including the closest entries
1526
+ at or before `start_time` and at or after `end_time` if applicable.
1527
+
1528
+ Warns
1529
+ -----
1530
+ UserWarning
1531
+ If no records at or before `start_time` or no records at or after `end_time` are found.
1532
+
1533
+ UserWarning
1534
+ If the dataset does not contain any time dimension or the time dimension is incorrectly named.
1535
+ """
1536
+
1537
+ time_dim = self.dim_names["time"]
1538
+
1539
+ ds = _select_relevant_times(ds, time_dim, self.start_time, self.end_time, False)
1540
+
1541
+ return ds
1542
+
1543
+ def compute_climatology(self):
1544
+ logging.info("Compute climatology for river forcing.")
1545
+
1546
+ time_dim = self.dim_names["time"]
1547
+
1548
+ flux = self.ds[self.var_names["flux"]].groupby(f"{time_dim}.month").mean()
1549
+ self.ds[self.var_names["flux"]] = flux
1550
+
1551
+ ds = assign_dates_to_climatology(self.ds, "month")
1552
+ ds = ds.swap_dims({"month": "time"})
1553
+ object.__setattr__(self, "ds", ds)
1554
+
1555
+ updated_dim_names = {**self.dim_names}
1556
+ updated_dim_names["time"] = "time"
1557
+ object.__setattr__(self, "dim_names", updated_dim_names)
1558
+
1559
+ object.__setattr__(self, "climatology", True)
1560
+
1561
+ def sort_by_river_volume(self, ds: xr.Dataset) -> xr.Dataset:
1562
+ """Sorts the dataset by river volume in descending order (largest rivers first),
1563
+ if the volume variable is available.
1564
+
1565
+ This method uses the river volume to reorder the dataset such that the rivers with
1566
+ the largest volumes come first in the `station` dimension. If the volume variable
1567
+ is not present in the dataset, a warning is logged.
1568
+
1569
+ Parameters
1570
+ ----------
1571
+ ds : xr.Dataset
1572
+ The xarray Dataset containing the river data to be sorted by volume.
1573
+
1574
+ Returns
1575
+ -------
1576
+ xr.Dataset
1577
+ The dataset with rivers sorted by their volume in descending order.
1578
+ If the volume variable is not available, the original dataset is returned.
1579
+ """
1580
+
1581
+ if "vol" in self.opt_var_names:
1582
+ volume_values = ds[self.opt_var_names["vol"]].values
1583
+ if isinstance(volume_values, np.ndarray):
1584
+ # Check if all volume values are the same
1585
+ if np.all(volume_values == volume_values[0]):
1586
+ # If all volumes are the same, no need to reverse order
1587
+ sorted_indices = np.argsort(
1588
+ volume_values
1589
+ ) # Sort in ascending order
1590
+ else:
1591
+ # If volumes differ, reverse order for descending sort
1592
+ sorted_indices = np.argsort(volume_values)[
1593
+ ::-1
1594
+ ] # Reverse for descending order
1595
+
1596
+ ds = ds.isel(**{self.dim_names["station"]: sorted_indices})
1597
+
1598
+ else:
1599
+ logging.warning("The volume data is not in a valid array format.")
1600
+ else:
1601
+ logging.warning(
1602
+ "Cannot sort rivers by volume. 'vol' is missing in the variable names."
1603
+ )
1604
+
1605
+ return ds
1606
+
1607
+ def extract_relevant_rivers(self, target_coords, dx):
1608
+ """Extracts a subset of the dataset based on the proximity of river mouths to
1609
+ target coordinates.
1610
+
1611
+ This method calculates the distance between each river mouth and the provided target coordinates
1612
+ (latitude and longitude) using the `gc_dist` function. It then filters the dataset to include only those
1613
+ river stations whose minimum distance from the target is less than a specified threshold distance (`dx`).
1614
+
1615
+ Parameters
1616
+ ----------
1617
+ target_coords : dict
1618
+ A dictionary containing the target coordinates for the comparison. It should include:
1619
+ - "lon" (float): The target longitude in degrees.
1620
+ - "lat" (float): The target latitude in degrees.
1621
+ - "straddle" (bool): A flag indicating whether to adjust the longitudes for stations that cross the
1622
+ International Date Line. If `True`, longitudes greater than 180 degrees are adjusted by subtracting 360,
1623
+ otherwise, negative longitudes are adjusted by adding 360.
1624
+
1625
+ dx : float
1626
+ The maximum distance threshold (in meters) for including a river station. Only river mouths that are
1627
+ within `dx` meters from the target coordinates will be included in the returned dataset.
1628
+
1629
+ Returns
1630
+ -------
1631
+ indices : dict
1632
+ A dictionary containing the indices of the rivers that are within the threshold distance from
1633
+ the target coordinates. The dictionary keys are:
1634
+ - "station" : numpy.ndarray
1635
+ The indices of the rivers that satisfy the distance threshold.
1636
+ - "eta_rho" : numpy.ndarray
1637
+ The indices of the `eta_rho` dimension corresponding to the selected stations.
1638
+ - "xi_rho" : numpy.ndarray
1639
+ The indices of the `xi_rho` dimension corresponding to the selected stations.
1640
+ """
1641
+
1642
+ # Retrieve longitude and latitude of river mouths
1643
+ river_lon = self.ds[self.var_names["longitude"]]
1644
+ river_lat = self.ds[self.var_names["latitude"]]
1645
+
1646
+ # Adjust longitude based on whether it crosses the International Date Line (straddle case)
1647
+ if target_coords["straddle"]:
1648
+ river_lon = xr.where(river_lon > 180, river_lon - 360, river_lon)
1649
+ else:
1650
+ river_lon = xr.where(river_lon < 0, river_lon + 360, river_lon)
1651
+
1652
+ # Calculate the distance between the target coordinates and each river mouth
1653
+ dist = gc_dist(target_coords["lon"], target_coords["lat"], river_lon, river_lat)
1654
+ dist_min = dist.min(dim=["eta_rho", "xi_rho"])
1655
+ # Filter the dataset to include only stations within the distance threshold
1656
+ if (dist_min < dx).any():
1657
+ ds = self.ds.where(dist_min < dx, drop=True)
1658
+ ds = self.sort_by_river_volume(ds)
1659
+ dist = dist.where(dist_min < dx, drop=True).transpose(
1660
+ self.dim_names["station"], "eta_rho", "xi_rho"
1661
+ )
1662
+ dist_min = dist_min.where(dist_min < dx, drop=True)
1663
+
1664
+ # Find the indices of the closest grid cell to the river mouth
1665
+ indices = np.where(dist == dist_min)
1666
+ names = (
1667
+ self.ds[self.var_names["name"]]
1668
+ .isel({self.dim_names["station"]: indices[0]})
1669
+ .values
1670
+ )
1671
+ # Return the indices in a dictionary format
1672
+ indices = {
1673
+ "station": indices[0],
1674
+ "eta_rho": indices[1],
1675
+ "xi_rho": indices[2],
1676
+ "name": names,
1677
+ }
1678
+ else:
1679
+ ds = xr.Dataset()
1680
+ indices = {
1681
+ "station": [],
1682
+ "eta_rho": [],
1683
+ "xi_rho": [],
1684
+ "name": [],
1685
+ }
1686
+
1687
+ object.__setattr__(self, "ds", ds)
1688
+
1689
+ return indices
1690
+
1691
+
1692
+ @dataclass(frozen=True, kw_only=True)
1693
+ class DaiRiverDataset(RiverDataset):
1694
+ """Represents river data from the Dai river dataset.
1695
+
1696
+ Parameters
1697
+ ----------
1698
+ filename : Union[str, Path, List[Union[str, Path]]], optional
1699
+ The path to the Dai River dataset file. If not provided, the dataset will be downloaded
1700
+ automatically via the `pooch` library.
1701
+ start_time : datetime
1702
+ The start time for selecting relevant data.
1703
+ end_time : datetime
1704
+ The end time for selecting relevant data.
1705
+ dim_names: Dict[str, str], optional
1706
+ Dictionary specifying the names of dimensions in the dataset.
1707
+ var_names: Dict[str, str], optional
1708
+ Dictionary of variable names that are required in the dataset.
1709
+ opt_var_names: Dict[str, str], optional
1710
+ Dictionary of variable names that are optional in the dataset.
1711
+ climatology : bool
1712
+ Indicates whether the dataset is climatological. Defaults to False.
1713
+
1714
+ Attributes
1715
+ ----------
1716
+ ds : xr.Dataset
1717
+ The xarray Dataset containing the forcing data on its original grid.
1718
+ """
1719
+
1720
+ filename: Union[str, Path, List[Union[str, Path]]] = field(
1721
+ default_factory=lambda: download_river_data("dai_trenberth_may2019.nc")
1722
+ )
1723
+ start_time: datetime
1724
+ end_time: datetime
1725
+ dim_names: Dict[str, str] = field(
1726
+ default_factory=lambda: {
1727
+ "station": "station",
1728
+ "time": "time",
1729
+ }
1730
+ )
1731
+ var_names: Dict[str, str] = field(
1732
+ default_factory=lambda: {
1733
+ "latitude": "lat_mou",
1734
+ "longitude": "lon_mou",
1735
+ "flux": "FLOW",
1736
+ "ratio": "ratio_m2s",
1737
+ "name": "riv_name",
1738
+ }
1739
+ )
1740
+ opt_var_names: Dict[str, str] = field(
1741
+ default_factory=lambda: {
1742
+ "vol": "vol_stn",
1743
+ }
1744
+ )
1745
+ climatology: Optional[bool] = False
1746
+ ds: xr.Dataset = field(init=False, repr=False)
1747
+
1748
+ def add_time_info(self, ds: xr.Dataset) -> xr.Dataset:
1749
+ """Adds time information to the dataset based on the climatology flag and
1750
+ dimension names.
1751
+
1752
+ This method processes the dataset to include time information according to the climatology
1753
+ setting. If the dataset represents climatology data and the time dimension is labeled as
1754
+ "month", it assigns dates to the dataset based on a monthly climatology. Additionally, it
1755
+ handles dimension name updates if necessary.
1756
+
1757
+ Parameters
1758
+ ----------
1759
+ ds : xr.Dataset
1760
+ The input dataset to which time information will be added.
1761
+
1762
+ Returns
1763
+ -------
1764
+ xr.Dataset
1765
+ The dataset with time information added, including adjustments for climatology and
1766
+ dimension names.
1767
+ """
1768
+ time_dim = self.dim_names["time"]
1769
+
1770
+ # Extract the 'time' variable as a numpy array
1771
+ time_vals = ds[time_dim].values
1772
+
1773
+ # Handle rounding of the time values
1774
+ year = np.round(time_vals * 1e-2).astype(int)
1775
+ month = np.round((time_vals * 1e-2 - year) * 1e2).astype(int)
1776
+
1777
+ # Convert to datetime (assuming the day is always 15th for this example)
1778
+ dates = [datetime(year=i, month=m, day=15) for i, m in zip(year, month)]
1779
+
1780
+ ds[time_dim] = dates
1781
+
1782
+ return ds
1783
+
1784
+
1785
+ # shared functions
1786
+
1787
+
1788
+ def _load_data(filename, dim_names, use_dask, decode_times=True):
1789
+ """Load dataset from the specified file.
1790
+
1791
+ Parameters
1792
+ ----------
1793
+ filename : Union[str, Path, List[Union[str, Path]]]
1794
+ The path to the data file(s). Can be a single string (with or without wildcards), a single Path object,
1795
+ or a list of strings or Path objects containing multiple files.
1796
+ dim_names: Dict[str, str], optional
1797
+ Dictionary specifying the names of dimensions in the dataset.
1798
+ use_dask: bool
1799
+ Indicates whether to use dask for chunking. If True, data is loaded with dask; if False, data is loaded eagerly. Defaults to False.
1800
+ decode_times: bool, optional
1801
+ If True, decode times encoded in the standard NetCDF datetime format into datetime objects. Otherwise, leave them encoded as numbers.
1802
+ Defaults to True.
1803
+
1804
+ Returns
1805
+ -------
1806
+ ds : xr.Dataset
1807
+ The loaded xarray Dataset containing the forcing data.
1808
+
1809
+ Raises
1810
+ ------
1811
+ FileNotFoundError
1812
+ If the specified file does not exist.
1813
+ ValueError
1814
+ If a list of files is provided but dim_names["time"] is not available or use_dask=False.
1815
+ """
1816
+
1817
+ # Precompile the regex for matching wildcard characters
1818
+ wildcard_regex = re.compile(r"[\*\?\[\]]")
1819
+
1820
+ # Convert Path objects to strings
1821
+ if isinstance(filename, (str, Path)):
1822
+ filename_str = str(filename)
1823
+ elif isinstance(filename, list):
1824
+ filename_str = [str(f) for f in filename]
1825
+ else:
1826
+ raise ValueError("filename must be a string, Path, or a list of strings/Paths.")
1827
+ # Handle the case when filename is a string
1828
+ contains_wildcard = False
1829
+ if isinstance(filename_str, str):
1830
+ contains_wildcard = bool(wildcard_regex.search(filename_str))
1831
+ if contains_wildcard:
1832
+ matching_files = glob.glob(filename_str)
1833
+ if not matching_files:
1834
+ raise FileNotFoundError(
1835
+ f"No files found matching the pattern '{filename_str}'."
1836
+ )
1837
+ else:
1838
+ matching_files = [filename_str]
1839
+
1840
+ # Handle the case when filename is a list
1841
+ elif isinstance(filename_str, list):
1842
+ contains_wildcard = any(wildcard_regex.search(f) for f in filename_str)
1843
+ if contains_wildcard:
1844
+ matching_files = []
1845
+ for f in filename_str:
1846
+ files = glob.glob(f)
1847
+ if not files:
1848
+ raise FileNotFoundError(
1849
+ f"No files found matching the pattern '{f}'."
1850
+ )
1851
+ matching_files.extend(files)
1852
+ else:
1853
+ matching_files = filename_str
1854
+
1855
+ # Check if time dimension is available when multiple files are provided
1856
+ if isinstance(filename_str, list) and "time" not in dim_names:
1857
+ raise ValueError(
1858
+ "A list of files is provided, but time dimension is not available. "
1859
+ "A time dimension must be available to concatenate the files."
1860
+ )
1861
+
1862
+ # Determine the kwargs for combining datasets
1863
+ if contains_wildcard or len(matching_files) == 1:
1864
+ # If there is a wildcard or just one file, use by_coords
1865
+ kwargs = {"combine": "by_coords"}
1866
+ else:
1867
+ # Otherwise, use nested combine based on time
1868
+ kwargs = {"combine": "nested", "concat_dim": dim_names["time"]}
1869
+
1870
+ # Base kwargs used for dataset combination
1871
+ combine_kwargs = {
1872
+ "coords": "minimal",
1873
+ "compat": "override",
1874
+ "combine_attrs": "override",
1875
+ }
1876
+
1877
+ if use_dask:
1878
+
1879
+ chunks = {
1880
+ dim_names["latitude"]: -1,
1881
+ dim_names["longitude"]: -1,
1882
+ }
1883
+ if "depth" in dim_names:
1884
+ chunks[dim_names["depth"]] = -1
1885
+ if "time" in dim_names:
1886
+ chunks[dim_names["time"]] = 1
1887
+
1888
+ ds = xr.open_mfdataset(
1889
+ matching_files,
1890
+ decode_times=decode_times,
1891
+ chunks=chunks,
1892
+ **combine_kwargs,
1893
+ **kwargs,
1894
+ )
1895
+ else:
1896
+ ds_list = []
1897
+ for file in matching_files:
1898
+ ds = xr.open_dataset(file, decode_times=decode_times, chunks=None)
1899
+ ds_list.append(ds)
1900
+
1901
+ if kwargs["combine"] == "by_coords":
1902
+ ds = xr.combine_by_coords(ds_list, **combine_kwargs)
1903
+ elif kwargs["combine"] == "nested":
1904
+ ds = xr.combine_nested(
1905
+ ds_list, concat_dim=kwargs["concat_dim"], **combine_kwargs
1906
+ )
1907
+
1908
+ if "time" in dim_names and dim_names["time"] not in ds.dims:
1909
+ ds = ds.expand_dims(dim_names["time"])
1910
+
1911
+ return ds
1912
+
1913
+
1914
+ def _check_dataset(
1915
+ ds: xr.Dataset,
1916
+ dim_names: Dict[str, str],
1917
+ var_names: Dict[str, str],
1918
+ opt_var_names: Optional[Dict[str, str]] = None,
1919
+ ) -> None:
1920
+ """Check if the dataset contains the specified variables and dimensions.
1921
+
1922
+ Parameters
1923
+ ----------
1924
+ ds : xr.Dataset
1925
+ The xarray Dataset to check.
1926
+ dim_names: Dict[str, str], optional
1927
+ Dictionary specifying the names of dimensions in the dataset.
1928
+ var_names: Dict[str, str]
1929
+ Dictionary of variable names that are required in the dataset.
1930
+ opt_var_names : Optional[Dict[str, str]], optional
1931
+ Dictionary of optional variable names.
1932
+ These variables are not strictly required, and the function will not raise an error if they are missing.
1933
+ Default is None, meaning no optional variables are considered.
1934
+
1935
+
1936
+ Raises
1937
+ ------
1938
+ ValueError
1939
+ If the dataset does not contain the specified variables or dimensions.
1940
+ """
1941
+ missing_dims = [dim for dim in dim_names.values() if dim not in ds.dims]
1942
+ if missing_dims:
1943
+ raise ValueError(
1944
+ f"Dataset does not contain all required dimensions. The following dimensions are missing: {missing_dims}"
1945
+ )
1946
+
1947
+ missing_vars = [var for var in var_names.values() if var not in ds.data_vars]
1948
+ if missing_vars:
1949
+ raise ValueError(
1950
+ f"Dataset does not contain all required variables. The following variables are missing: {missing_vars}"
1951
+ )
1952
+
1953
+ if opt_var_names:
1954
+ missing_optional_vars = [
1955
+ var for var in opt_var_names.values() if var not in ds.data_vars
1956
+ ]
1957
+ if missing_optional_vars:
1958
+ logging.warning(
1959
+ f"Optional variables missing (but not critical): {missing_optional_vars}"
1960
+ )
1961
+
1962
+
1963
+ def _select_relevant_times(
1964
+ ds, time_dim, start_time=None, end_time=None, climatology=False
1965
+ ) -> xr.Dataset:
1966
+ """Select a subset of the dataset based on the specified time range.
1967
+
1968
+ This method filters the dataset to include all records between `start_time` and `end_time`.
1969
+ Additionally, it ensures that one record at or before `start_time` and one record at or
1970
+ after `end_time` are included, even if they fall outside the strict time range.
1971
+
1972
+ If no `end_time` is specified, the method will select the time range of
1973
+ [start_time, start_time + 24 hours] and return the closest time entry to `start_time` within that range.
1974
+
1975
+ Parameters
1976
+ ----------
1977
+ ds : xr.Dataset
1978
+ The input dataset to be filtered. Must contain a time dimension.
1979
+ time_dim: str
1980
+ Name of time dimension.
1981
+ start_time : Optional[datetime], optional
1982
+ The start time for selecting relevant data. If not provided, the data is not filtered by start time.
1983
+ end_time : Optional[datetime], optional
1984
+ The end time for selecting relevant data. If not provided, only data at the start_time is selected if start_time is provided,
1985
+ or no filtering is applied if start_time is not provided.
1986
+ climatology : bool
1987
+ Indicates whether the dataset is climatological. Defaults to False.
1988
+
1989
+ Returns
1990
+ -------
1991
+ xr.Dataset
1992
+ A dataset filtered to the specified time range, including the closest entries
1993
+ at or before `start_time` and at or after `end_time` if applicable.
1994
+
1995
+ Raises
1996
+ ------
1997
+ ValueError
1998
+ If no matching times are found between `start_time` and `start_time + 24 hours`.
1999
+
2000
+ Warns
2001
+ -----
2002
+ UserWarning
2003
+ If the dataset contains exactly 12 time steps but the climatology flag is not set.
2004
+ This may indicate that the dataset represents climatology data.
2005
+
2006
+ UserWarning
2007
+ If no records at or before `start_time` or no records at or after `end_time` are found.
2008
+
2009
+ UserWarning
2010
+ If the dataset does not contain any time dimension or the time dimension is incorrectly named.
2011
+
2012
+ Notes
2013
+ -----
2014
+ - If the `climatology` flag is set and `end_time` is not provided, the method will
2015
+ interpolate initial conditions from climatology data.
2016
+ - If the dataset uses `cftime` datetime objects, these will be converted to standard
2017
+ `np.datetime64` objects before filtering.
2018
+ """
2019
+
2020
+ if time_dim in ds.variables:
2021
+ if climatology:
2022
+ if len(ds[time_dim]) != 12:
2023
+ raise ValueError(
2024
+ f"The dataset contains {len(ds[time_dim])} time steps, but the climatology flag is set to True, which requires exactly 12 time steps."
2025
+ )
2026
+ if not end_time:
2027
+ # Interpolate from climatology for initial conditions
2028
+ ds = interpolate_from_climatology(ds, time_dim, start_time)
2029
+ else:
2030
+ time_type = get_time_type(ds[time_dim])
2031
+ if time_type == "int":
2032
+ raise ValueError(
2033
+ "The dataset contains integer time values, which are only supported when the climatology flag is set to True. However, your climatology flag is set to False."
2034
+ )
2035
+ if time_type == "cftime":
2036
+ ds = ds.assign_coords(
2037
+ {time_dim: convert_cftime_to_datetime(ds[time_dim])}
2038
+ )
2039
+ if end_time:
2040
+ end_time = end_time
2041
+
2042
+ # Identify records before or at start_time
2043
+ before_start = ds[time_dim] <= np.datetime64(start_time)
2044
+ if before_start.any():
2045
+ closest_before_start = (
2046
+ ds[time_dim].where(before_start, drop=True).max()
2047
+ )
2048
+ else:
2049
+ logging.warning("No records found at or before the start_time.")
2050
+ closest_before_start = ds[time_dim].min()
2051
+
2052
+ # Identify records after or at end_time
2053
+ after_end = ds[time_dim] >= np.datetime64(end_time)
2054
+ if after_end.any():
2055
+ closest_after_end = ds[time_dim].where(after_end, drop=True).min()
2056
+ else:
2057
+ logging.warning("No records found at or after the end_time.")
2058
+ closest_after_end = ds[time_dim].max()
2059
+
2060
+ # Select records within the time range and add the closest before/after
2061
+ within_range = (ds[time_dim] > np.datetime64(start_time)) & (
2062
+ ds[time_dim] < np.datetime64(end_time)
2063
+ )
2064
+ selected_times = ds[time_dim].where(
2065
+ within_range
2066
+ | (ds[time_dim] == closest_before_start)
2067
+ | (ds[time_dim] == closest_after_end),
2068
+ drop=True,
2069
+ )
2070
+ ds = ds.sel({time_dim: selected_times})
2071
+ else:
2072
+ # Look in time range [start_time, start_time + 24h]
2073
+ end_time = start_time + timedelta(days=1)
2074
+ times = (np.datetime64(start_time) <= ds[time_dim]) & (
2075
+ ds[time_dim] < np.datetime64(end_time)
2076
+ )
2077
+ if np.all(~times):
2078
+ raise ValueError(
2079
+ f"The dataset does not contain any time entries between the specified start_time: {start_time} "
2080
+ f"and {start_time + timedelta(hours=24)}. "
2081
+ "Please ensure the dataset includes time entries for that range."
2082
+ )
2083
+
2084
+ ds = ds.where(times, drop=True)
2085
+ if ds.sizes[time_dim] > 1:
2086
+ # Pick the time closest to start_time
2087
+ ds = ds.isel({time_dim: 0})
2088
+ logging.info(
2089
+ f"Selected time entry closest to the specified start_time ({start_time}) within the range [{start_time}, {start_time + timedelta(hours=24)}]: {ds[time_dim].values}"
2090
+ )
2091
+ else:
2092
+ logging.warning(
2093
+ "Dataset does not contain any time information. Please check if the time dimension "
2094
+ "is correctly named or if the dataset includes time data."
2095
+ )
2096
+
2097
+ return ds
2098
+
2099
+
2100
+ def decode_string(byte_array):
2101
+
2102
+ # Decode each byte and handle errors with 'ignore'
2103
+ decoded_string = "".join(
2104
+ [
2105
+ x.decode("utf-8", errors="ignore") # Ignore invalid byte sequences
2106
+ for x in byte_array.values
2107
+ if isinstance(x, bytes) and x != b" " and x is not np.nan
2108
+ ]
2109
+ )
2110
+
2111
+ return decoded_string