Coverage for /home/runner/work/nr-catalog-tools/nr-catalog-tools/nrcatalogtools/sxs.py: 81%

1import os 

2 

3import sxs 

4from nrcatalogtools import catalog, waveform 

5 

6 

7class SXSCatalog(catalog.CatalogBase): 

8 def __init__(self, catalog=None, verbosity=0, **kwargs) -> None: 

9 super().__init__(catalog, **kwargs) 

10 self._verbosity = verbosity 

11 self._add_paths_to_metadata() 

12 

13 def waveform_filename_from_simname(self, sim_name): 

14 return os.path.basename(self.waveform_filepath_from_simname(sim_name)) 

15 

16 def waveform_filepath_from_simname(self, sim_name): 

17 poss_files = self.select_files(f"{sim_name}/Lev/rhOverM") 

18 file_path = ( 

19 sxs.sxs_directory("cache") 

20 / poss_files[list(poss_files.keys())[0]]["truepath"] 

21 ) 

22 if not os.path.exists(file_path): 

23 if self._verbosity > 2: 

24 print( 

25 f"WARNING: Could not resolve path for {sim_name}" 

26 f"..best calculated path = {file_path}" 

27 ) 

28 return file_path.as_posix() 

29 

30 def metadata_filename_from_simname(self, sim_name): 

31 return os.path.basename(self.metadata_filepath_from_simname(sim_name)) 

32 

33 def metadata_filepath_from_simname(self, sim_name): 

34 poss_files = self.select_files(f"{sim_name}/Lev/metadata.json") 

35 file_path = ( 

36 sxs.sxs_directory("cache") 

37 / poss_files[list(poss_files.keys())[0]]["truepath"] 

38 ) 

39 if not os.path.exists(file_path): 

40 if self._verbosity > 2: 

41 print( 

42 f"WARNING: Could not resolve path for {sim_name}" 

43 f"..best calculated path = {file_path}" 

44 ) 

45 return file_path.as_posix() 

46 

47 def get(self, sim_name, extrapolation_order=2, download=None): 

48 extrap_key = f"Extrapolated_N{extrapolation_order}.dir" 

49 

50 # Download only if not available 

51 if download is None: 

52 try: 

53 raw_obj = sxs.load(f"{sim_name}/Lev/rhOverM", download=False) 

54 except Exception as excep: 

55 excep("Waveform data not available. Setting download to True") 

56 raw_obj = sxs.load(f"{sim_name}/Lev/rhOverM", download=True) 

57 else: 

58 raw_obj = sxs.load(f"{sim_name}/Lev/rhOverM", download=download) 

59 

60 raw_obj = raw_obj.get(extrap_key) 

61 # Get the sim metadata 

62 sim_metadata = self.get_metadata(sim_name) 

63 # Get the h5 file name 

64 filepath = self.waveform_filepath_from_simname(sim_name) 

65 # Add to these details to the metadata 

66 raw_obj._metadata.update( 

67 {"metadata": sim_metadata, "waveform_data_location": filepath} 

68 ) 

69 return waveform.WaveformModes( 

70 raw_obj.data, sim_metadata=sim_metadata, **raw_obj._metadata 

71 ) 

72 

73 def download_waveform_data(self, sim_name): 

74 _ = sxs.load(f"{sim_name}/Lev/metadata.json") 

75 return sxs.load(f"{sim_name}/Lev/rhOverM") 

76 

77 def waveform_url_from_simname(self, sim_name): 

78 raise NotImplementedError("This shouldn't be called.") 

79 

80 def _add_paths_to_metadata(self): 

81 metadata_dict = self._dict["simulations"] 

82 existing_cols = list(metadata_dict[list(metadata_dict.keys())[0]].keys()) 

83 new_cols = [ 

84 "metadata_link", 

85 "metadata_location", 

86 "waveform_data_link", 

87 "waveform_data_location", 

88 ] 

89 

90 if any([col not in existing_cols for col in new_cols]): 

91 for sim_name in metadata_dict: 

92 if "metadata_location" not in existing_cols: 

93 metadata_dict[sim_name][ 

94 "metadata_location" 

95 ] = self.metadata_filepath_from_simname(sim_name) 

96 if "metadata_link" not in existing_cols: 

97 metadata_dict[sim_name]["metadata_link"] = "" 

98 if "waveform_data_link" not in existing_cols: 

99 metadata_dict[sim_name]["waveform_data_link"] = "" 

100 if "waveform_data_location" not in existing_cols: 

101 metadata_dict[sim_name][ 

102 "waveform_data_location" 

103 ] = self.waveform_filepath_from_simname(sim_name)