plotSolution.py 226 KB
Newer Older
Loic Hausammann's avatar
Loic Hausammann committed
1
#!/usr/bin/env python3
lhausamm's avatar
lhausamm committed
2
3
4
5
6
7
#######################################################################
#
#  UNIFIED ANALYSIS SCRIPT FOR DISK SIMULATION FOR THE AGORA PROJECT
#
#  FOR SCRIPT HISTORY SEE VERSION CONTROL CHANGELOG
#
Loic Hausammann's avatar
Loic Hausammann committed
8
#  Note: This script works with yt-3.5.1.
lhausamm's avatar
lhausamm committed
9
10
#
#######################################################################
lhausamm's avatar
lhausamm committed
11
# This script is a copy of the AGORA project (https://bitbucket.org/mornkr/agora-analysis-script/)
lhausamm's avatar
lhausamm committed
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
# modified in order to take into account SWIFT
import matplotlib
matplotlib.use('Agg')
import sys
import math
import copy
import csv
import numpy as np
import matplotlib.colorbar as cb
import matplotlib.lines as ln
import yt.utilities.physical_constants as constants
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from yt.mods import *
from yt.units.yt_array import YTQuantity
from mpl_toolkits.axes_grid1 import AxesGrid
Loic Hausammann's avatar
Loic Hausammann committed
28
from matplotlib.offsetbox import AnchoredText
lhausamm's avatar
lhausamm committed
29
30
31
32
33
34
35
36
from matplotlib.ticker import MaxNLocator
from yt.fields.particle_fields import add_volume_weighted_smoothed_field
from yt.fields.particle_fields import add_nearest_neighbor_field
from yt.analysis_modules.star_analysis.api import StarFormationRate
from yt.analysis_modules.halo_analysis.api import *
from yt.data_objects.particle_filters import add_particle_filter
from scipy.stats import kde
from subprocess import call
Loic Hausammann's avatar
Gear    
Loic Hausammann committed
37
# mylog.setLevel(1)
lhausamm's avatar
lhausamm committed
38

39
40
from yt.utilities.math_utils import get_cyl_r, get_cyl_theta, get_cyl_z

41
import sys
lhausamm's avatar
lhausamm committed
42

lhausamm's avatar
lhausamm committed
43
44
draw_density_map                 = 1#1           # 0/1         = OFF/ON
draw_temperature_map             = 1#1           # 0/1         = OFF/ON
lhausamm's avatar
lhausamm committed
45
draw_cellsize_map                = 2#2           # 0/1/2/3     = OFF/ON/ON now with resolution map where particle_size is defined as [M/rho]^(1/3) for SPH/ON with both cell_size and resolution maps
lhausamm's avatar
lhausamm committed
46
47
48
draw_elevation_map               = 1#1           # 0/1         = OFF/ON
draw_metal_map                   = 0#0           # 0/1/2       = OFF/ON/ON with total metal mass in the simulation annotated (this will be inaccurate for SPH)
draw_zvel_map                    = 0#0           # 0/1         = OFF/ON
49
50
51
52
53
54
55
draw_star_map                    = 1#1           # 0/1         = OFF/ON
draw_star_clump_stats            = 1#1           # 0/1/2/3     = OFF/ON/ON with additional star_map with annotated clumps/ON with addtional star_map + extra clump mass function with GIZMO-ps2
draw_SFR_map                     = 2###2         # 0/1/2       = OFF/ON/ON with local K-S plot using patches
draw_PDF                         = 2###2         # 0/1/2/3     = OFF/ON/ON with constant pressure/entropy lines/ON with additional annotations such as 1D profile from a specific code (CHANGA)
draw_pos_vel_PDF                 = 4##4          # 0/1/2/3/4   = OFF/ON/ON with 1D profile/ON also with 1D dispersion profile/ON also with separate 1D vertical dispersion profiles
draw_star_pos_vel_PDF            = 4##4          # 0/1/2/3/4   = OFF/ON/ON with 1D profile/ON also with 1D dispersion profile/ON also with separate 1D vertical dispersion profiles
draw_rad_height_PDF              = 2##2          # 0/1/2/3     = OFF/ON/ON with 1D profile/ON with analytic ftn subtracted
Loic Hausammann's avatar
Gear    
Loic Hausammann committed
56
draw_metal_PDF                   = 1##1          # 0/1         = OFF/ON
57
58
59
60
61
62
draw_density_DF                  = 2#2           # 0/1/2       = OFF/ON/ON with difference plot between 1st and 2nd datasets (when 2, dataset_num should be set to 2)
draw_radius_DF                   = 1#1           # 0/1         = OFF/ON
draw_star_radius_DF              = 2#2           # 0/1/2       = OFF/ON/ON with SFR profile and K-S plot (when 2, this automatically turns on draw_radius_DF)
draw_height_DF                   = 1#1           # 0/1         = OFF/ON
draw_SFR                         = 1#1           # 0/1/2       = OFF/ON/ON with extra line with GIZMO-ps2
draw_cut_through                 = 0#0           # 0/1         = OFF/ON
lhausamm's avatar
lhausamm committed
63
add_nametag                      = 1#1           # 0/1         = OFF/ON
64
add_mean_fractional_dispersion   = 1#1           # 0/1         = OFF/ON
lhausamm's avatar
lhausamm committed
65

66
dataset_num                      = 2             # 1/2         = 1st dataset(Grackle+noSF)/2nd dataset(Grackle+SF+ThermalFbck) for AGORA Paper 4
lhausamm's avatar
lhausamm committed
67
68
69
yt_version_pre_3_2_3             = 0             # 0/1         = NO/YES to "Is the yt version being used pre yt-dev-3.2.3?"
times                            = [0, 500]      # in Myr
figure_width                     = 30            # in kpc
lhausamm's avatar
lhausamm committed
70
71
n_ref                            = 64            # 8; for SPH codes
over_refine_factor               = 1             # 2; for SPH codes
lhausamm's avatar
lhausamm committed
72
73
74
75
76
77
78
79
80
aperture_size_SFR_map            = 750           # in pc       = Used if draw_SFR_map = 1, 750 matches Bigiel et al. (2008)
young_star_cutoff_SFR_map        = 20            # in Myr      = Used if draw_SFR_map = 1
young_star_cutoff_star_radius_DF = 20            # in Myr      = Used if draw_star_radius_DF = 1
mean_dispersion_radius_range     = [2, 10]       # in kpc      = Used if add_mean_fractional_dispersion = 1, for draw_pos_vel_PDF, draw_star_pos_vel_PDF, draw_rad_height_PDF, draw_radius_DF, draw_star_radius_DF
mean_dispersion_height_range     = [0, 0.6]      # in kpc      = Used if add_mean_fractional_dispersion = 1, for draw_height_DF
mean_dispersion_time_range       = [50, 500]     # in Myr      = Used if add_mean_fractional_dispersion = 1, for draw_SFR
mean_dispersion_density_range    = [1e-25, 1e-22]# in g/cm3    = Used if add_mean_fractional_dispersion = 1, for draw_density_DF
disk_normal_vector               = [0., 0., 1.]

81
82
gadget_default_unit_base = {'UnitLength_in_cm'         : 3.08567758e+21,
                            'UnitMass_in_g'            :    1.98848e+43,
83
                            'UnitVelocity_in_cm_per_s' :      100000}
lhausamm's avatar
lhausamm committed
84
85
86
87
88
color_names              = ['red', 'magenta', 'orange', 'gold', 'green', 'cyan', 'blue', 'blueviolet', 'black']
linestyle_names          = ['-']
marker_names             = ['s', 'o', 'p', 'v', '^', '<', '>', 'h', '*']

# file_location = ['/lustre/ki/pfs/mornkr/080112_CHaRGe/pfs-hyades/AGORA-DISK-repository-for-use/Grackle+noSF/',
89
#                '/lustre/ki/pfs/mornkr/080112_CHaRGe/pfs-hyades/AGORA-DISK-repository-for-use/Grackle+SF+ThermalFbck/']
lhausamm's avatar
lhausamm committed
90
# file_location = ['/global/project/projectdirs/agora/AGORA-DISK-repository-for-use/Grackle+noSF/',
91
#                '/global/project/projectdirs/agora/AGORA-DISK-repository-for-use/Grackle+SF+ThermalFbck/']
lhausamm's avatar
lhausamm committed
92
93
# codes = ['ART-I', 'ART-II', 'ENZO', 'RAMSES', 'CHANGA', 'GASOLINE', 'GADGET-3', 'GEAR', 'GIZMO']
# filenames = [[[file_location[0]+'ART-I/IC/AGORA_Galaxy_LOW.d', file_location[0]+'ART-I/t0.5Gyr/10MpcBox_csf512_02350.d'],
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
#             [file_location[0]+'ART-II/noSF_def_2p/OUT/AGORA_LOW_000000.art', file_location[0]+'ART-II/noSF_def_2p/OUT/AGORA_LOW_000098.art'],
#             [file_location[0]+'ENZO/DD0000/DD0000', file_location[0]+'ENZO/DD0100/DD0100'],
#             [file_location[0]+'RAMSES/output_00001/info_00001.txt', file_location[0]+'RAMSES/output_00068/info_00068.txt'],
#             [file_location[0]+'CHANGA/disklow/disklow.000000', file_location[0]+'CHANGA/disklow/disklow.000500'],
#             [file_location[0]+'GASOLINE/LOW_dataset1.00001', file_location[0]+'GASOLINE/LOW_dataset1.00335'],
#             [file_location[0]+'GADGET-3/AGORA_ISO_LOW_DRY/snap_iso_dry_000.hdf5', file_location[0]+'GADGET-3/AGORA_ISO_LOW_DRY/snap_iso_dry_010.hdf5'],
#             [file_location[0]+'GEAR/snapshot_0000', file_location[0]+'GEAR/snapshot_0500'],
#             [file_location[0]+'GIZMO/snapshot_temp_000', file_location[0]+'GIZMO/snapshot_temp_100']],
#            [[file_location[1]+'ART-I/IC/AGORA_Galaxy_LOW.d', file_location[1]+'ART-I/t0.5Gyr/10MpcBox_csf512_02350.d'],
#             [file_location[1]+'ART-II/SF_FBth_def_2p/OUT/AGORA_LOW_000000.art', file_location[1]+'ART-II/SF_FBth_def_2p/OUT/AGORA_LOW_000311.art'],
#             [file_location[1]+'ENZO/DD0000/DD0000', file_location[1]+'ENZO/DD0050/DD0050'],
#             [file_location[1]+'RAMSES/output_00001/info_00001.txt', file_location[1]+'RAMSES/output_00216/info_00216.txt'],
#             [file_location[1]+'CHANGA/disklow/disklow.000000', file_location[1]+'CHANGA/disklow/disklow.000500'],
#             [file_location[1]+'GASOLINE/LOW_dataset1.00001', file_location[1]+'GASOLINE/LOW_dataset2.00335'],
#             [file_location[1]+'GADGET-3/AGORA_ISO_LOW_SF_SNII_Thermal_Chevalier_SFT10/snap_iso_sf_000.hdf5', file_location[1]+'GADGET-3/AGORA_ISO_LOW_SF_SNII_Thermal_Chevalier_SFT10/snap_iso_sf_010.hdf5'],
#             [file_location[1]+'GEAR/snapshot_0000', file_location[1]+'GEAR/snapshot_0500'],
#             [file_location[1]+'GIZMO/snapshot_temp_000', file_location[1]+'GIZMO/snapshot_temp_100']]]
111
codes = ['SWIFT', 'GEAR']
Loic Hausammann's avatar
Loic Hausammann committed
112
filenames = [[["./agora_disk_IC.hdf5", "./agora_disk_500Myr.hdf5"], # Sim-noSFF
Loic Hausammann's avatar
Gear    
Loic Hausammann committed
113
              ["./snapshot_0000.hdf5", "./snapshot_0050.hdf5"]],
Loic Hausammann's avatar
Loic Hausammann committed
114
             [["./agora_disk_IC.hdf5", "./agora_disk_500Myr.hdf5"], # Sim-SFF (with star formation and feedback)
Loic Hausammann's avatar
Gear    
Loic Hausammann committed
115
              ["./snapshot_0000.hdf5", "./snapshot_0050.hdf5"]]] # I did not check the order, they can be switched
116
117
118
119

# codes = ["SWIFT"]
# filenames = [[["./agora_disk_0000.hdf5", "./agora_disk_0050.hdf5"]],
#             [["./agora_disk_0000.hdf5", "./agora_disk_0050.hdf5"]]]
lhausamm's avatar
lhausamm committed
120
121
# codes = ['ART-I']
# filenames = [[[file_location[0]+'ART-I/IC/AGORA_Galaxy_LOW.d', file_location[0]+'ART-I/t0.5Gyr/10MpcBox_csf512_02350.d']],
122
#            [[file_location[1]+'ART-I/IC/AGORA_Galaxy_LOW.d', file_location[1]+'ART-I/t0.5Gyr/10MpcBox_csf512_02350.d']]]
lhausamm's avatar
lhausamm committed
123
124
# codes = ['ART-II']
# filenames = [[[file_location[0]+'ART-II/noSF_def_2p/OUT/AGORA_LOW_000000.art', file_location[0]+'ART-II/noSF_def_2p/OUT/AGORA_LOW_000098.art']],
125
#            [[file_location[1]+'ART-II/SF_FBth_def_2p/OUT/AGORA_LOW_000000.art', file_location[1]+'ART-II/SF_FBth_def_2p/OUT/AGORA_LOW_000311.art']]]
lhausamm's avatar
lhausamm committed
126
127
# codes = ['ENZO']
# filenames = [[[file_location[0]+'ENZO/DD0000/DD0000', file_location[0]+'ENZO/DD0100/DD0100']],
128
#            [[file_location[1]+'ENZO/DD0000/DD0000', file_location[1]+'ENZO/DD0050/DD0050']]]
lhausamm's avatar
lhausamm committed
129
130
# codes = ['RAMSES']
# filenames = [[[file_location[0]+'RAMSES/output_00001/info_00001.txt', file_location[0]+'RAMSES/output_00068/info_00068.txt']],
131
#            [[file_location[1]+'RAMSES/output_00001/info_00001.txt', file_location[1]+'RAMSES/output_00216/info_00216.txt']]]
lhausamm's avatar
lhausamm committed
132
# codes = ['CHANGA']
133
134
# filenames = [[[file_location[0]+'CHANGA/disklow/disklow.000000', file_location[0]+'CHANGA/disklow/disklow.000500']],
#            [[file_location[1]+'CHANGA/disklow/disklow.000000', file_location[1]+'CHANGA/disklow/disklow.000500']]]
lhausamm's avatar
lhausamm committed
135
136
# codes = ['GASOLINE']
# filenames = [[[file_location[0]+'GASOLINE/LOW_dataset1.00001', file_location[0]+'GASOLINE/LOW_dataset1.00335']],
137
#            [[file_location[1]+'GASOLINE/LOW_dataset1.00001', file_location[1]+'GASOLINE/LOW_dataset2.00335']]]
lhausamm's avatar
lhausamm committed
138
139
# codes = ['GADGET-3']
# filenames = [[[file_location[0]+'GADGET-3/AGORA_ISO_LOW_DRY/snap_iso_dry_000.hdf5', file_location[0]+'GADGET-3/AGORA_ISO_LOW_DRY/snap_iso_dry_010.hdf5']],
140
#            [[file_location[1]+'GADGET-3/AGORA_ISO_LOW_SF_SNII_Thermal_Chevalier_SFT10/snap_iso_sf_000.hdf5', file_location[1]+'GADGET-3/AGORA_ISO_LOW_SF_SNII_Thermal_Chevalier_SFT10/snap_iso_sf_010.hdf5']]]
lhausamm's avatar
lhausamm committed
141
# codes = ['GEAR']
Loic Hausammann's avatar
Gear    
Loic Hausammann committed
142
143
# filenames = [[['snapshot_0000.hdf5', 'snapshot_0500.hdf5']],
#            [['snapshot_0000.hdf5', 'snapshot_0500']]]
lhausamm's avatar
lhausamm committed
144
145
# codes = ['GIZMO']
# filenames = [[[file_location[0]+'GIZMO/snapshot_temp_000', file_location[0]+'GIZMO/snapshot_temp_100']],
146
#            [[file_location[1]+'GIZMO/snapshot_temp_000', file_location[1]+'GIZMO/snapshot_temp_100']]]
lhausamm's avatar
lhausamm committed
147
148

def load_dataset(dataset_num, time, code, codes, filenames_entry):
149
150
151
152
153
154
155
156
157
158
159
        if codes[code] == 'ART-I': # ART frontend doesn't find the accompanying files, so we specify them; see http://yt-project.org/docs/dev/examining/loading_data.html#art-data
                dirnames = filenames_entry[code][time][:filenames_entry[code][time].rfind('/')+1]
                if time == 0:
                        timestamp = ''
                else:
                        timestamp = filenames_entry[code][time][filenames_entry[code][time].rfind('_'):filenames_entry[code][time].rfind('.')]
                pf = load(filenames_entry[code][time], file_particle_header=dirnames+'PMcrd'+timestamp+'.DAT', file_particle_data=dirnames+'PMcrs0'+timestamp+'.DAT', file_particle_stars=dirnames+'stars'+timestamp+'.dat')
        elif codes[code] == 'CHANGA' or codes[code] == 'GASOLINE': # For TIPSY frontend, always make sure to place your parameter file in the same directory as your datasets
                pf = load(filenames_entry[code][time], n_ref=n_ref, over_refine_factor=over_refine_factor)
        elif codes[code] == 'GADGET-3': # For GADGET-3 2nd dataset, there somehow exist particles very far from the center; so we use [-2000, 2000] for a bounding_box
                pf = load(filenames_entry[code][time], unit_base = gadget_default_unit_base, bounding_box=[[-2000.0, 2000.0], [-2000.0, 2000.0], [-2000.0, 2000.0]], n_ref=n_ref, over_refine_factor=over_refine_factor)
lhausamm's avatar
lhausamm committed
160
        elif codes[code] == "SWIFT":
161
162
                pf = load(filenames_entry[code][time], unit_base = gadget_default_unit_base, bounding_box=[[-2000.0, 2000.0], [-2000.0, 2000.0], [-2000.0, 2000.0]], n_ref=n_ref, over_refine_factor=over_refine_factor)
        elif codes[code] == 'GEAR':
Loic Hausammann's avatar
Loic Hausammann committed
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
                pf = load(filenames_entry[code][time], unit_base = gadget_default_unit_base, bounding_box=[[-2000.0, 2000.0], [-2000.0, 2000.0], [-2000.0, 2000.0]], n_ref=n_ref, over_refine_factor=over_refine_factor) 
               # from yt.frontends.gadget.definitions import gadget_header_specs
               #  from yt.frontends.gadget.definitions import gadget_ptype_specs
               #  from yt.frontends.gadget.definitions import gadget_field_specs
               #  if with_cooling:
               #          gadget_header_specs["chemistry"] = (('h1',  4, 'c'),('h2',  4, 'c'),('empty',  256, 'c'),)
               #          header_spec = "default+chemistry"
               #  else:
               #          header_spec = "default"
               #  gear_ptype_specs = ("Gas", "Stars", "Halo", "Disk", "Bulge", "Bndry")
               #  if dataset_num == 1:
               #          pf = GadgetDataset(filenames_entry[code][time], unit_base = gadget_default_unit_base, bounding_box=[[-1000.0, 1000.0], [-1000.0, 1000.0], [-1000.0, 1000.0]], header_spec=header_spec,
               #                             ptype_spec=gear_ptype_specs, n_ref=n_ref, over_refine_factor=over_refine_factor)
               #  elif dataset_num == 2:
               #          # For GEAR 2nd dataset (Grackle+SF+ThermalFbck), "Metals" is acutally 10-species field; check how metallicity field is added below.
               #          if with_cooling:
               #                  agora_gear  = ( "Coordinates",
               #                                  "Velocities",
               #                                  "ParticleIDs",
               #                                  "Mass",
               #                                  ("InternalEnergy", "Gas"),
               #                                  ("Density", "Gas"),
               #                                  ("SmoothingLength", "Gas"),
               #                                  ("Metals", "Gas"),
               #                                  ("StellarFormationTime", "Stars"),
               #                                  ("StellarInitMass", "Stars"),
               #                                  ("StellarIDs", "Stars"),
               #                                  ("StellarDensity", "Stars"),
               #                                  ("StellarSmoothingLength", "Stars"),
               #                                  ("StellarMetals", "Stars"),
               #                                  ("Opt1", "Stars"),
               #                                  ("Opt2", "Stars"),
               #                  )
               #          else:
               #                  agora_gear  = ( "Coordinates",
               #                                  "Velocities",
               #                                  "ParticleIDs",
               #                                  "Mass",
               #                                  ("InternalEnergy", "Gas"),
               #                                  ("Density", "Gas"),
               #                                  ("SmoothingLength", "Gas"),
               #                  )
               #          gadget_field_specs["agora_gear"] = agora_gear
               #          pf = GadgetDataset(filenames_entry[code][time], unit_base = gadget_default_unit_base, bounding_box=[[-1000.0, 1000.0], [-1000.0, 1000.0], [-1000.0, 1000.0]], header_spec=header_spec,
               #                             ptype_spec=gear_ptype_specs, field_spec="agora_gear", n_ref=n_ref, over_refine_factor=over_refine_factor)
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
        elif codes[code] == 'GIZMO':
                from yt.frontends.gadget.definitions import gadget_field_specs
                if dataset_num == 1:
                        agora_gizmo = ( "Coordinates",
                                        "Velocities",
                                        "ParticleIDs",
                                        "Mass",
                                        ("Temperature", "Gas"),
                                        ("Density", "Gas"),
                                        ("Electron_Number_Density", "Gas"),
                                        ("HI_NumberDensity", "Gas"),
                                        ("SmoothingLength", "Gas"),
                                        )
                elif dataset_num == 2:
                        agora_gizmo = ( "Coordinates",
                                        "Velocities",
                                        "ParticleIDs",
                                        "Mass",
                                        ("Temperature", "Gas"),
                                        ("Density", "Gas"),
                                        ("ElectronAbundance", "Gas"),
                                        ("NeutralHydrogenAbundance", "Gas"),
                                        ("SmoothingLength", "Gas"),
                                        ("StarFormationRate", "Gas"),
                                        ("StellarFormationTime", "Stars"),
                                        ("Metallicity", ("Gas", "Stars")),
                                        ("DelayTime", "Stars"),
                                        ("StellarInitMass", "Stars"),
                                        )
                gadget_field_specs["agora_gizmo"] = agora_gizmo
                pf = load(filenames_entry[code][time], unit_base = gadget_default_unit_base, bounding_box=[[-1000.0, 1000.0], [-1000.0, 1000.0], [-1000.0,1000.0]], field_spec="agora_gizmo",
                          n_ref=n_ref, over_refine_factor=over_refine_factor)
        else:
                pf = load(filenames_entry[code][time])
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
                                # Bug with some fields, need to redefine them in order to get correct units
        def _radius(field, data):
            """The cylindrical radius component of the particle positions
                            
            Relative to the coordinate system defined by the *normal* vector
            and *center* field parameters.
            """
            normal = data.get_field_parameter('normal')
            pos = data['relative_particle_position'].T
            return data.ds.arr(get_cyl_r(pos, normal), 'cm')

        pf.add_field(("PartType0", "radius"),
                     sampling_type="particle",
                     function=_radius,
                     units="cm",
                     validators=[ValidateParameter("normal"), ValidateParameter("center")])
        pf.add_field(("PartType4", "radius"),
                     sampling_type="particle",
                     function=_radius,
                     units="cm",
                     validators=[ValidateParameter("normal"), ValidateParameter("center")])
        pf.add_field(("PartType1", "radius"),
                     sampling_type="particle",
                     function=_radius,
                     units="cm",
                     validators=[ValidateParameter("normal"), ValidateParameter("center")])
        pf.add_field(("all", "radius"),
                     sampling_type="particle",
                     function=_radius,
                     units="cm",
                     validators=[ValidateParameter("normal"), ValidateParameter("center")])

        def _height(field, data):
            """The cylindrical radius component of the particle positions
            
            Relative to the coordinate system defined by the *normal* vector
            and *center* field parameters.
            """
            normal = data.get_field_parameter('normal')
            pos = data['relative_particle_position'].T
            return data.ds.arr(get_cyl_z(pos, normal), "cm")

        pf.add_field(("PartType0", "height"),
                     sampling_type="particle",
                     function=_height,
                     units="cm",
                     validators=[ValidateParameter("normal"), ValidateParameter("center")])
        pf.add_field(("PartType1", "height"),
                     sampling_type="particle",
                     function=_height,
                     units="cm",
                     validators=[ValidateParameter("normal"), ValidateParameter("center")])
        pf.add_field(("PartType4", "height"),
                     sampling_type="particle",
                     function=_height,
                     units="cm",
                     validators=[ValidateParameter("normal"), ValidateParameter("center")])
        pf.add_field(("all", "height"),
                     sampling_type="particle",
                     function=_height,
                     units="cm",
                     validators=[ValidateParameter("normal"), ValidateParameter("center")])

305
306
307
        return pf

fig_density_map             = []
lhausamm's avatar
lhausamm committed
308
309
310
311
fig_temperature_map         = []
fig_cellsize_map            = []
fig_cellsize_map_2          = []
fig_elevation_map           = []
312
313
314
315
316
fig_metal_map               = []
fig_zvel_map                = []
fig_star_map                = []
fig_star_map_2              = []
fig_star_map_3              = []
lhausamm's avatar
lhausamm committed
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
fig_degr_sfr_map            = []
fig_degr_density_map        = []
fig_PDF                     = []
fig_pos_vel_PDF             = []
fig_star_pos_vel_PDF        = []
fig_rad_height_PDF          = []
fig_metal_PDF               = []
grid_density_map            = []
grid_temperature_map        = []
grid_cellsize_map           = []
grid_cellsize_map_2         = []
grid_elevation_map          = []
grid_metal_map              = []
grid_zvel_map               = []
grid_star_map               = []
grid_star_map_2             = []
grid_star_map_3             = []
grid_degr_sfr_map           = []
grid_degr_density_map       = []
grid_PDF                    = []
grid_pos_vel_PDF            = []
grid_star_pos_vel_PDF       = []
grid_rad_height_PDF         = []
grid_metal_PDF              = []
star_clump_masses_hop       = []
star_clump_masses_fof       = []
star_clump_masses_hop_ref   = []
star_clump_masses_fof_ref   = []
pos_vel_xs                  = []
pos_vel_profiles            = []
pos_disp_xs                 = []
pos_disp_profiles           = []
pos_disp_vert_xs            = []
pos_disp_vert_profiles      = []
star_pos_vel_xs             = []
star_pos_vel_profiles       = []
star_pos_disp_xs            = []
star_pos_disp_profiles      = []
star_pos_disp_vert_xs       = []
star_pos_disp_vert_profiles = []
rad_height_xs               = []
rad_height_profiles         = []
density_DF_xs               = []
density_DF_profiles         = []
density_DF_1st_xs           = []
density_DF_1st_profiles     = []
radius_DF_xs                = []
radius_DF_profiles          = []
surface_density             = []
star_radius_DF_xs           = []
star_radius_DF_profiles     = []
star_surface_density        = []
sfr_radius_DF_xs            = []
sfr_radius_DF_profiles      = []
sfr_surface_density         = []
height_DF_xs                = []
height_DF_profiles          = []
height_surface_density      = []
sfr_ts                      = []
sfr_cum_masses              = []
sfr_sfrs                    = []
surf_dens_SFR_map           = []
sfr_surf_dens_SFR_map       = []
cut_through_zs              = []
cut_through_zvalues         = []
cut_through_xs              = []
cut_through_xvalues         = []

for time in range(len(times)):
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
        if draw_density_map == 1:
                fig_density_map      += [plt.figure(figsize=(100,20))]
                grid_density_map     += [AxesGrid(fig_density_map[time], (0.01,0.01,0.99,0.99), nrows_ncols = (2, len(codes)), axes_pad = 0.02, add_all = True, share_all = True,
                                                  label_mode = "1", cbar_mode = "single", cbar_location = "right", cbar_size = "2%", cbar_pad = 0.02)]
        if draw_temperature_map == 1:
                fig_temperature_map  += [plt.figure(figsize=(100,20))]
                grid_temperature_map += [AxesGrid(fig_temperature_map[time], (0.01,0.01,0.99,0.99), nrows_ncols = (2, len(codes)), axes_pad = 0.02, add_all = True, share_all = True,
                                                  label_mode = "1", cbar_mode = "single", cbar_location = "right", cbar_size = "2%", cbar_pad = 0.02)]
        if draw_cellsize_map >= 1:
                fig_cellsize_map     += [plt.figure(figsize=(100,20))]
                grid_cellsize_map    += [AxesGrid(fig_cellsize_map[time], (0.01,0.01,0.99,0.99), nrows_ncols = (2, len(codes)), axes_pad = 0.02, add_all = True, share_all = True,
                                                  label_mode = "1", cbar_mode = "single", cbar_location = "right", cbar_size = "2%", cbar_pad = 0.02)]
                fig_cellsize_map_2   += [plt.figure(figsize=(100,20))]
                grid_cellsize_map_2  += [AxesGrid(fig_cellsize_map_2[time], (0.01,0.01,0.99,0.99), nrows_ncols = (2, len(codes)), axes_pad = 0.02, add_all = True, share_all = True,
                                                  label_mode = "1", cbar_mode = "single", cbar_location = "right", cbar_size = "2%", cbar_pad = 0.02)]
        if draw_elevation_map == 1:
                fig_elevation_map    += [plt.figure(figsize=(100,20))]
                grid_elevation_map   += [AxesGrid(fig_elevation_map[time], (0.01,0.01,0.99,0.99), nrows_ncols = (1, len(codes)), axes_pad = 0.02, add_all = True, share_all = True,
                                                  label_mode = "1", cbar_mode = "single", cbar_location = "right", cbar_size = "6%", cbar_pad = 0.02)]
        if draw_metal_map >= 1:
                fig_metal_map        += [plt.figure(figsize=(100,20))]
                grid_metal_map       += [AxesGrid(fig_metal_map[time], (0.01,0.01,0.99,0.99), nrows_ncols = (2, len(codes)), axes_pad = 0.02, add_all = True, share_all = True,
                                                  label_mode = "1", cbar_mode = "single", cbar_location = "right", cbar_size = "2%", cbar_pad = 0.02)]
        if draw_zvel_map == 1:
                fig_zvel_map         += [plt.figure(figsize=(100,20))]
                grid_zvel_map        += [AxesGrid(fig_zvel_map[time], (0.01,0.01,0.99,0.99), nrows_ncols = (2, len(codes)), axes_pad = 0.02, add_all = True, share_all = True,
                                                  label_mode = "1", cbar_mode = "single", cbar_location = "right", cbar_size = "2%", cbar_pad = 0.02)]
        if draw_star_map == 1:
                fig_star_map         += [plt.figure(figsize=(100,20))]
                grid_star_map        += [AxesGrid(fig_star_map[time], (0.01,0.01,0.99,0.99), nrows_ncols = (2, len(codes)), axes_pad = 0.02, add_all = True, share_all = True,
                                                  label_mode = "1", cbar_mode = "single", cbar_location = "right", cbar_size = "2%", cbar_pad = 0.02)]
        if draw_star_clump_stats >= 1:
                star_clump_masses_hop.append([])
                star_clump_masses_fof.append([])
                fig_star_map_2       += [plt.figure(figsize=(100,20))]
                grid_star_map_2      += [AxesGrid(fig_star_map_2[time], (0.01,0.01,0.99,0.99), nrows_ncols = (2, len(codes)), axes_pad = 0.02, add_all = True, share_all = True,
                                                  label_mode = "1", cbar_mode = "single", cbar_location = "right", cbar_size = "2%", cbar_pad = 0.02)]
                fig_star_map_3       += [plt.figure(figsize=(100,20))]
                grid_star_map_3      += [AxesGrid(fig_star_map_3[time], (0.01,0.01,0.99,0.99), nrows_ncols = (2, len(codes)), axes_pad = 0.02, add_all = True, share_all = True,
                                                  label_mode = "1", cbar_mode = "single", cbar_location = "right", cbar_size = "2%", cbar_pad = 0.02)]
        if draw_SFR_map >= 1:
                fig_degr_density_map += [plt.figure(figsize=(100,20))]
                grid_degr_density_map+= [AxesGrid(fig_degr_density_map[time], (0.01,0.01,0.99,0.99), nrows_ncols = (1, len(codes)), axes_pad = 0.02, add_all = True, share_all = True,
                                                  label_mode = "1", cbar_mode = "single", cbar_location = "right", cbar_size = "6%", cbar_pad = 0.02)]
                fig_degr_sfr_map     += [plt.figure(figsize=(100,20))]
                grid_degr_sfr_map    += [AxesGrid(fig_degr_sfr_map[time], (0.01,0.01,0.99,0.99), nrows_ncols = (1, len(codes)), axes_pad = 0.02, add_all = True, share_all = True,
                                                  label_mode = "1", cbar_mode = "single", cbar_location = "right", cbar_size = "6%", cbar_pad = 0.02)]
                surf_dens_SFR_map.append([])
                sfr_surf_dens_SFR_map.append([])
        if draw_SFR_map >= 2 or draw_star_radius_DF >= 2:
                def import_text(filename, separator):
                        for line in csv.reader(open(filename), delimiter=separator, skipinitialspace=True):
                                if line:
                                        yield line
        if draw_PDF >= 1:
                fig_PDF              += [plt.figure(figsize=(50, 80))]
Loic Hausammann's avatar
Gear    
Loic Hausammann committed
442
                grid_PDF             += [AxesGrid(fig_PDF[time], (0.01,0.01,0.99,0.99), nrows_ncols = (2, int(math.ceil(len(codes)/2.0))), axes_pad = 0.05, add_all = True, share_all = True,
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
                                                  label_mode = "1", cbar_mode = "single", cbar_location = "right", cbar_size = "2%", cbar_pad = 0.05, aspect = False)]
        if draw_pos_vel_PDF >= 1:
                fig_pos_vel_PDF      += [plt.figure(figsize=(50, 80))]
                grid_pos_vel_PDF     += [AxesGrid(fig_pos_vel_PDF[time], (0.01,0.01,0.99,0.99), nrows_ncols = (3, int(math.ceil(len(codes)/3.0))), axes_pad = 0.05, add_all = True, share_all = True,
                                                  label_mode = "1", cbar_mode = "single", cbar_location = "right", cbar_size = "2%", cbar_pad = 0.05, aspect = False)]
                pos_vel_xs.append([])
                pos_vel_profiles.append([])
                pos_disp_xs.append([])
                pos_disp_profiles.append([])
                pos_disp_vert_xs.append([])
                pos_disp_vert_profiles.append([])
        if draw_star_pos_vel_PDF >= 1:
                fig_star_pos_vel_PDF += [plt.figure(figsize=(50, 80))]
                grid_star_pos_vel_PDF+= [AxesGrid(fig_star_pos_vel_PDF[time], (0.01,0.01,0.99,0.99), nrows_ncols = (3, int(math.ceil(len(codes)/3.0))), axes_pad = 0.05, add_all = True, share_all = True,
                                                  label_mode = "1", cbar_mode = "single", cbar_location = "right", cbar_size = "2%", cbar_pad = 0.05, aspect = False)]
                star_pos_vel_xs.append([])
                star_pos_vel_profiles.append([])
                star_pos_disp_xs.append([])
                star_pos_disp_profiles.append([])
                star_pos_disp_vert_xs.append([])
                star_pos_disp_vert_profiles.append([])
        if draw_rad_height_PDF >= 1:
                fig_rad_height_PDF   += [plt.figure(figsize=(50, 80))]
                grid_rad_height_PDF  += [AxesGrid(fig_rad_height_PDF[time], (0.01,0.01,0.99,0.99), nrows_ncols = (3, int(math.ceil(len(codes)/3.0))), axes_pad = 0.05, add_all = True, share_all = True,
                                                  label_mode = "1", cbar_mode = "single", cbar_location = "right", cbar_size = "2%", cbar_pad = 0.05, aspect = False)]
                rad_height_xs.append([])
                rad_height_profiles.append([])
        if draw_metal_PDF == 1:
                fig_metal_PDF        += [plt.figure(figsize=(50, 80))]
Loic Hausammann's avatar
Gear    
Loic Hausammann committed
472
                grid_metal_PDF       += [AxesGrid(fig_metal_PDF[time], (0.01,0.01,0.99,0.99), nrows_ncols = (2, int(math.ceil(len(codes)/2.0))), axes_pad = 0.05, add_all = True, share_all = True,
473
474
475
476
477
478
479
                                                  label_mode = "1", cbar_mode = "single", cbar_location = "right", cbar_size = "2%", cbar_pad = 0.05, aspect = False)]
        if draw_density_DF >= 1:
                density_DF_xs.append([])
                density_DF_profiles.append([])
                density_DF_1st_xs.append([])
                density_DF_1st_profiles.append([])
                if draw_density_DF == 2 and dataset_num == 1:
Loic Hausammann's avatar
Loic Hausammann committed
480
                        print("This won't work; resetting draw_density_DF to 1...")
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
                        draw_density_DF == 1
        if draw_star_radius_DF >= 1:
                star_radius_DF_xs.append([])
                star_radius_DF_profiles.append([])
                star_surface_density.append([])
                sfr_radius_DF_xs.append([])
                sfr_radius_DF_profiles.append([])
                sfr_surface_density.append([])
        if draw_star_radius_DF == 2:
                draw_radius_DF = 1
        if draw_radius_DF == 1:
                radius_DF_xs.append([])
                radius_DF_profiles.append([])
                surface_density.append([])
        if draw_height_DF == 1:
                height_DF_xs.append([])
                height_DF_profiles.append([])
                height_surface_density.append([])
        if draw_SFR >= 1:
                sfr_ts.append([])
                sfr_cum_masses.append([])
                sfr_sfrs.append([])
        if draw_cut_through == 1:
                cut_through_zs.append([])
                cut_through_zvalues.append([])
                cut_through_xs.append([])
                cut_through_xvalues.append([])
lhausamm's avatar
lhausamm committed
508
509
510

for time in range(len(times)):

511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
        for code in range(len(codes)):

                if (dataset_num != 1 and dataset_num != 2) or (filenames[dataset_num-1][code] == []):
                        continue

                ####################################
                #        PRE-ANALYSIS STEPS        #
                ####################################

                # LOAD DATASETS
                pf = load_dataset(dataset_num, time, code, codes, filenames[dataset_num-1])

                # PARTICLE FILED NAMES FOR SPH CODES, AND STELLAR PARTICLE FILTERS
                PartType_Gas_to_use = "Gas"
                PartType_Star_to_use = "Stars"
                PartType_StarBeforeFiltered_to_use = "Stars"
                MassType_to_use = "Mass"
                MetallicityType_to_use = "Metallicity"
                FormationTimeType_to_use = "StellarFormationTime" # for GADGET/GEAR/GIZMO, this field has to be added in frontends/sph/fields.py, in which only "FormationTime" can be recognized
                SmoothingLengthType_to_use = "SmoothingLength"

                if codes[code] == 'CHANGA' or codes[code] == 'GASOLINE':
                        MetallicityType_to_use = "Metals"
                        PartType_Star_to_use = "NewStars"
                        PartType_StarBeforeFiltered_to_use = "Stars"
                        FormationTimeType_to_use = "FormationTime"
                        SmoothingLengthType_to_use = "smoothing_length"
                        def NewStars(pfilter, data): # see http://yt-project.org/docs/dev/analyzing/filtering.html#filtering-particle-fields
                                return (data[(pfilter.filtered_type, FormationTimeType_to_use)] > 0)
                        add_particle_filter(PartType_Star_to_use, function=NewStars, filtered_type=PartType_StarBeforeFiltered_to_use, requires=[FormationTimeType_to_use])
                        pf.add_particle_filter(PartType_Star_to_use)
                        pf.periodicity = (True, True, True) # this is needed especially when bPeriodic = 0 in GASOLINE, to avoid RuntimeError in geometry/selection_routines.pyx:855
                elif codes[code] == 'ART-I':
                        PartType_StarBeforeFiltered_to_use = "stars"
                        FormationTimeType_to_use = "particle_creation_time"
                        def Stars(pfilter, data):
                                return (data[(pfilter.filtered_type, FormationTimeType_to_use)] > 0)
#                               return ((data[(pfilter.filtered_type, FormationTimeType_to_use)] > 0) & (data[(pfilter.filtered_type, "particle_index")] >= 212500)) # without the fix metioned above, the above line won't work because all "stars"="specie1" have the same wrong particle_creation_time of 6.851 Gyr; so you will have to use this quick and dirty patch
                        add_particle_filter(PartType_Star_to_use, function=Stars, filtered_type=PartType_StarBeforeFiltered_to_use, requires=[FormationTimeType_to_use, "particle_index"])
                        pf.add_particle_filter(PartType_Star_to_use)
                elif codes[code] == 'ART-II':
                        PartType_StarBeforeFiltered_to_use = "STAR"
                        if time != 0: # BIRTH_TIME field exists in ART-II but as a dimensionless quantity for some reason in frontends/artio/fields.py; so we create StellarFormationTime field
                                def _FormationTime(field, data):
                                        return pf.arr(data["STAR", "BIRTH_TIME"].d, 'code_time')
                                pf.add_field(("STAR", FormationTimeType_to_use), function=_FormationTime, particle_type=True, take_log=False, units="code_time")
                        def Stars(pfilter, data):
                                return (data[(pfilter.filtered_type, FormationTimeType_to_use)] > 0)
                        add_particle_filter(PartType_Star_to_use, function=Stars, filtered_type=PartType_StarBeforeFiltered_to_use, requires=[FormationTimeType_to_use])
                        pf.add_particle_filter(PartType_Star_to_use)
                elif codes[code] == 'ENZO':
                        PartType_StarBeforeFiltered_to_use = "all"
                        FormationTimeType_to_use = "creation_time"
                        def Stars(pfilter, data):
                                return ((data[(pfilter.filtered_type, "particle_type")] == 2) & (data[(pfilter.filtered_type, FormationTimeType_to_use)] > 0))
                        add_particle_filter(PartType_Star_to_use, function=Stars, filtered_type=PartType_StarBeforeFiltered_to_use, requires=["particle_type", FormationTimeType_to_use])
                        pf.add_particle_filter(PartType_Star_to_use)
                elif codes[code] == "GADGET-3":
                        PartType_Gas_to_use = "PartType0"
                        PartType_Star_to_use = "PartType4"
                        PartType_StarBeforeFiltered_to_use = "PartType4"
                        MassType_to_use = "Masses"
                elif codes[code] == "SWIFT":
                        PartType_Gas_to_use = "PartType0"
Loic Hausammann's avatar
Loic Hausammann committed
575
576
                        PartType_Star_to_use = "PartType4"
                        PartType_StarBeforeFiltered_to_use = "PartType4"
577
578
                        if time != 0:
                            def _FormationTime(field, data):
Loic Hausammann's avatar
Gear    
Loic Hausammann committed
579
                                return pf.arr(data["PartType4", "BirthTimes"].d, 'code_time')
580
                            pf.add_field(("PartType4", FormationTimeType_to_use), function=_FormationTime, particle_type=True, take_log=False, units="code_time")
Loic Hausammann's avatar
Gear    
Loic Hausammann committed
581

582
                        MassType_to_use = "Masses"
Loic Hausammann's avatar
Loic Hausammann committed
583
584
585
586
587
                elif codes[code] == "GEAR":
                        PartType_Gas_to_use = "PartType0"
                        PartType_Star_to_use = "PartType1"
                        PartType_StarBeforeFiltered_to_use = "PartType1"
                        MassType_to_use = "Masses"
588
589
590
591
                        if time != 0:
                            def _FormationTime(field, data):
                                return pf.arr(data["PartType1", "StarFormationTime"].d, 'code_time')
                            pf.add_field(("PartType1", FormationTimeType_to_use), function=_FormationTime, particle_type=True, take_log=False, units="code_time")
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
                elif codes[code] == 'RAMSES':
                        PartType_StarBeforeFiltered_to_use = "all"
                        pf.current_time = pf.arr(pf.parameters['time'], 'code_time') # reset pf.current_time because it is incorrectly set up in frontends/ramses/data_structure.py, and I don't wish to mess with units there
                        FormationTimeType_to_use = "particle_age" # particle_age field actually means particle creation time, at least for this particular dataset, so the new field below is not needed
                        # if time != 0: # Only particle_age field exists in RAMSES (only for new stars + IC stars), so we create StellarFormationTime field
                        #       def _FormationTime(field, data):
                        #               return pf.current_time - data["all", "particle_age"].in_units("s")
                        #       pf.add_field(("all", FormationTimeType_to_use), function=_FormationTime, particle_type=True, take_log=False, units="code_time")
                        def Stars(pfilter, data):
                                return (data[(pfilter.filtered_type, "particle_age")] > 0)
                        add_particle_filter(PartType_Star_to_use, function=Stars, filtered_type=PartType_StarBeforeFiltered_to_use, requires=["particle_age"])
                        pf.add_particle_filter(PartType_Star_to_use)

                # AXIS SWAP FOR PLOT COLLECTION
                pf.coordinates.x_axis[1] = 0
                pf.coordinates.y_axis[1] = 2
                pf.coordinates.x_axis['y'] = 0
                pf.coordinates.y_axis['y'] = 2

                # ADDITIONAL FIELDS I: GLOBALLY USED FIELDS
                def _density_squared(field, data):
                        return data[("gas", "density")]**2
                pf.add_field(("gas", "density_squared"), function=_density_squared, units="g**2/cm**6")

                # ADDITIONAL FIELDS II: FOR CELL SIZE AND RESOLUTION MAPS
                def _CellSizepc(field,data):
                        return (data[("index", "cell_volume")].in_units('pc**3'))**(1/3.)
                pf.add_field(("index", "cell_size"), function=_CellSizepc, units='pc', display_name="Resolution $\Delta$ x", take_log=True )
                def _Inv2CellVolumeCode(field,data):
                        return data[("index", "cell_volume")]**-2
                pf.add_field(("index", "cell_volume_inv2"), function=_Inv2CellVolumeCode, units='code_length**(-6)', display_name="Inv2CellVolumeCode", take_log=True)
                if draw_cellsize_map >= 2:
                        if codes[code] == 'CHANGA' or codes[code] == 'GEAR' or codes[code] == 'GADGET-3' or codes[code] == 'GASOLINE' or codes[code] == 'GIZMO' or codes[code] == "SWIFT":
                                def _ParticleSizepc(field, data):
                                        return (data[(PartType_Gas_to_use, MassType_to_use)]/data[(PartType_Gas_to_use, "Density")])**(1./3.)
                                pf.add_field((PartType_Gas_to_use, "particle_size"), function=_ParticleSizepc, units="pc", display_name="Resolution $\Delta$ x", particle_type=True, take_log=True)
                                def _Inv2ParticleVolumepc(field, data):
                                        return (data[(PartType_Gas_to_use, MassType_to_use)]/data[(PartType_Gas_to_use, "Density")])**(-2.)
                                pf.add_field((PartType_Gas_to_use, "particle_volume_inv2"), function=_Inv2ParticleVolumepc, units="pc**(-6)", display_name="Inv2ParticleVolumepc", particle_type=True, take_log=True)
                                # Also creating smoothed field following an example in yt-project.org/docs/dev/cookbook/calculating_information.html; use hardcoded num_neighbors as in frontends/gadget/fields.py
                                fn = add_volume_weighted_smoothed_field(PartType_Gas_to_use, "Coordinates", MassType_to_use, SmoothingLengthType_to_use, "Density", "particle_size", pf.field_info, nneighbors=64)
                                fn = add_volume_weighted_smoothed_field(PartType_Gas_to_use, "Coordinates", MassType_to_use, SmoothingLengthType_to_use, "Density", "particle_volume_inv2", pf.field_info, nneighbors=64)

                                # Alias doesn't work -- e.g. pf.field_info.alias(("gas", "particle_size"), fn[0]) -- check alias below; so I simply add ("gas", "particle_size")
                                def _ParticleSizepc_2(field, data):
                                        return data["deposit", PartType_Gas_to_use+"_smoothed_"+"particle_size"]
                                pf.add_field(("gas", "particle_size"), function=_ParticleSizepc_2, units="pc", force_override=True, display_name="Resolution $\Delta$ x", particle_type=False, take_log=True)
                                def _Inv2ParticleVolumepc_2(field, data):
                                        return data["deposit", PartType_Gas_to_use+"_smoothed_"+"particle_volume_inv2"]
                                pf.add_field(("gas", "particle_volume_inv2"), function=_Inv2ParticleVolumepc_2, units="pc**(-6)", force_override=True, display_name="Inv2ParticleVolumepc", particle_type=False, take_log=True)

                # ADDITIONAL FIELDS III: TEMPERATURE
                if codes[code] == 'GEAR' or codes[code] == 'GADGET-3' or codes[code] == 'RAMSES' or codes[code] == "SWIFT":
                        # From grackle/src/python/utilities/convenience.py: Calculate a tabulated approximation to mean molecular weight (valid for data that used Grackle 2.0 or below)
                        def calc_mu_table_local(temperature):
                                tt = np.array([1.0e+01, 1.0e+02, 1.0e+03, 1.0e+04, 1.3e+04, 2.1e+04, 3.4e+04, 6.3e+04, 1.0e+05, 1.0e+09])
                                mt = np.array([1.18701555, 1.15484424, 1.09603514, 0.9981496, 0.96346395, 0.65175895, 0.6142901, 0.6056833, 0.5897776, 0.58822635])
                                logttt= np.log(temperature)
                                logmu = np.interp(logttt,np.log(tt),np.log(mt)) # linear interpolation in log-log space
                                return np.exp(logmu)
                        temperature_values = []
                        mu_values = []
                        T_over_mu_values = []
                        current_temperature = 1e1
                        final_temperature = 1e7
                        dlogT = 0.1
                        while current_temperature < final_temperature:
                                temperature_values.append(current_temperature)
                                current_mu = calc_mu_table_local(current_temperature)
                                mu_values.append(current_mu)
                                T_over_mu_values.append(current_temperature/current_mu)
                                current_temperature = np.exp(np.log(current_temperature)+dlogT)
                        def convert_T_over_mu_to_T(T_over_mu):
                                logT_over_mu = np.log(T_over_mu)
                                logT = np.interp(logT_over_mu, np.log(T_over_mu_values), np.log(temperature_values)) # linear interpolation in log-log space
                                return np.exp(logT)
                        if codes[code] == 'GEAR' or codes[code] == 'GADGET-3' or codes[code] == "SWIFT":
                                def _Temperature_3(field, data):
                                        gamma = 5.0/3.0
                                        T_over_mu = (data[PartType_Gas_to_use, "InternalEnergy"] * (gamma-1) * constants.mass_hydrogen_cgs / constants.boltzmann_constant_cgs).in_units('K').d # T/mu
                                        return YTArray(convert_T_over_mu_to_T(T_over_mu), 'K') # now T
                                pf.add_field((PartType_Gas_to_use, "Temperature"), function=_Temperature_3, particle_type=True, force_override=True, units="K")
                        elif codes[code] == 'RAMSES':
                                # The pressure field includes the artificial pressure support term, so one needs to be careful (compare with the exsiting frontends/ramses/fields.py)
                                def _temperature_3(field, data):
                                        T_J = 1800.0  # in K
                                        n_J = 8.0     # in H/cc
                                        gamma_0 = 2.0
                                        x_H = 0.76
                                        mH = 1.66e-24      # from pymses/utils/constants/__init__.py  (vs. in yt, mass_hydrogen_cgs = 1.007947*amu_cgs = 1.007947*1.660538921e-24 = 1.6737352e-24)
                                        kB = 1.3806504e-16 # from pymses/utils/constants/__init__.py  (vs. in yt, boltzmann_constant_cgs = 1.3806488e-16)
                                        if time != 0:
                                                T_over_mu = data["gas", "pressure"].d/data["gas", "density"].d * mH / kB - T_J * (data["gas", "density"].d * x_H / mH / n_J)**(gamma_0 - 1.0) # T/mu = T2 in Ramses
                                        else:
                                                T_over_mu = data["gas", "pressure"].d/data["gas", "density"].d * mH / kB # IC: no pressure support
                                        return YTArray(convert_T_over_mu_to_T(T_over_mu), 'K') # now T
                                pf.add_field(("gas", "temperature"), function=_temperature_3, force_override=True, units="K")

                # ADDITIONAL FIELDS IV: METALLICITY (IN MASS FRACTION, NOT IN ZSUN)
                if draw_metal_map >= 1 or draw_metal_PDF == 1:
                        if codes[code] == 'ART-I': # metallicity field in ART-I has a different meaning (see frontends/art/fields.py), and metallicity field in ART-II is missing
                                def _metallicity_2(field, data):
                                        return (data["gas", "metal_ii_density"] + data["gas", "metal_ia_density"]) / data["gas", "density"] # metal_ia_density needs to be added to account for initial 0.5 Zsun, even though we don't have SNe Ia
                                pf.add_field(("gas", "metallicity"), function=_metallicity_2, force_override=True, display_name="Metallicity", take_log=True, units="")
                        elif codes[code] == 'ART-II':
                                def _metallicity_2(field, data):
                                        return data["gas", "metal_ii_density"] / data["gas", "density"]
                                pf.add_field(("gas", "metallicity"), function=_metallicity_2, force_override=True, display_name="Metallicity", take_log=True, units="")
                        elif codes[code] == 'ENZO': # metallicity field in ENZO is in Zsun, so we create a new field
                                def _metallicity_2(field, data):
                                        return data["gas", "metal_density"] / data["gas", "density"]
                                pf.add_field(("gas", "metallicity"), function=_metallicity_2, force_override=True, display_name="Metallicity", take_log=True, units="")
                        elif codes[code] == 'GEAR': # "Metals" in GEAR is 10-species field ([:,9] is the total metal fraction), so requires a change in _vector_fields in frontends/gadget/io.py: added ("Metals", 10)
                                def _metallicity_2(field, data):
                                        if len(data[PartType_Gas_to_use, "Metals"].shape) == 1:
                                                return data[PartType_Gas_to_use, "Metals"]
                                        else:
                                                return data[PartType_Gas_to_use, "Metals"][:,9].in_units("") # in_units("") turned out to be crucial!; otherwise code_metallicity will be used and it will mess things up
                                # We are creating ("Gas", "Metallicity") here, different from ("Gas", "metallicity") which is auto-generated by yt but doesn't work properly
                                pf.add_field((PartType_Gas_to_use, MetallicityType_to_use), function=_metallicity_2, display_name="Metallicity", particle_type=True, take_log=True, units="")
                                # Also creating smoothed field following an example in yt-project.org/docs/dev/cookbook/calculating_information.html; use hardcoded num_neighbors as in frontends/gadget/fields.py
                                fn = add_volume_weighted_smoothed_field(PartType_Gas_to_use, "Coordinates", MassType_to_use, "SmoothingLength", "Density", MetallicityType_to_use, pf.field_info, nneighbors=64)
                                # Alias doesn't work -- e.g. pf.field_info.alias(("gas", "metallicity"), fn[0]) -- probably because pf=GadgetDataset()?, not load()?; so I add and replace existing ("gas", "metallicity")
                                def _metallicity_3(field, data):
                                        return data["deposit", PartType_Gas_to_use+"_smoothed_"+MetallicityType_to_use]
                                pf.add_field(("gas", "metallicity"), function=_metallicity_3, force_override=True, display_name="Metallicity", particle_type=False, take_log=True, units="")
Loic Hausammann's avatar
Gear    
Loic Hausammann committed
718
719
720
721
722
723
724
725
726
727
728
729
730
731
                        elif codes[code] == 'SWIFT': # "Metals" in SWIFT is 10-species field ([:,9] is the total metal fraction), so requires a change in _vector_fields in frontends/gadget/io.py: added ("Metals", 10)
                                def _metallicity_2(field, data):
                                        if len(data[PartType_Gas_to_use, "SmoothedElementAbundances"].shape) == 1:
                                                return data[PartType_Gas_to_use, "SmoothedElementAbundances"]
                                        else:
                                                return data[PartType_Gas_to_use, "SmoothedElementAbundances"][:,9].in_units("") # in_units("") turned out to be crucial!; otherwise code_metallicity will be used and it will mess things up
                                # We are creating ("Gas", "Metallicity") here, different from ("Gas", "metallicity") which is auto-generated by yt but doesn't work properly
                                pf.add_field((PartType_Gas_to_use, MetallicityType_to_use), function=_metallicity_2, display_name="Metallicity", particle_type=True, take_log=True, units="")
                                # Also creating smoothed field following an example in yt-project.org/docs/dev/cookbook/calculating_information.html; use hardcoded num_neighbors as in frontends/gadget/fields.py
                                fn = add_volume_weighted_smoothed_field(PartType_Gas_to_use, "Coordinates", MassType_to_use, "SmoothingLength", "Density", MetallicityType_to_use, pf.field_info, nneighbors=64)
                                # Alias doesn't work -- e.g. pf.field_info.alias(("gas", "metallicity"), fn[0]) -- probably because pf=GadgetDataset()?, not load()?; so I add and replace existing ("gas", "metallicity")
                                def _metallicity_3(field, data):
                                        return data["deposit", PartType_Gas_to_use+"_smoothed_"+MetallicityType_to_use]
                                pf.add_field(("gas", "metallicity"), function=_metallicity_3, force_override=True, display_name="Metallicity", particle_type=False, take_log=True, units="")
732
733
734
735
736
737
738
739
740
741
742
743
744
                        if draw_metal_map >= 2:
                                def _metal_mass(field, data):
                                        return data["gas", "cell_mass"] * data["gas", "metallicity"]
                                pf.add_field(("gas", "metal_mass"), function=_metal_mass, force_override=True, display_name="Metal Mass", take_log=True, units="Msun")

                # ADDITIONAL FIELDS V: FAKE PARTICLE FIELDS, CYLINDRICAL COORDINATES, etc.
                def rho_agora_disk(r, z):
                        r_d = YTArray(3.432, 'kpc')
                        z_d = 0.1*r_d
                        M_d = YTArray(4.297e10, 'Msun')
                        f_gas = 0.2
                        r_0 = (f_gas*M_d / (4.0*np.pi*r_d**2*z_d)).in_units('g/cm**3')
                        return r_0*numpy.exp(-r/r_d)*numpy.exp(-z/z_d)
lhausamm's avatar
lhausamm committed
745
746

                if codes[code] == "ART-I" or codes[code] == "ART-II" or codes[code] == "ENZO"  or codes[code] == "RAMSES":
747
748
749
750
751
752
753
754
                        def _cylindrical_z_abs(field, data):
                                return numpy.abs(data[("index", "cylindrical_z")])
                        pf.add_field(("index", "cylindrical_z_abs"), function=_cylindrical_z_abs, take_log=False, particle_type=False, units="cm")
                        def _density_minus_analytic(field, data):
                                return data[("gas", "density")] - rho_agora_disk(data[("index", "cylindrical_r")], data[("index", "cylindrical_z_abs")])
                        pf.add_field(("gas", "density_minus_analytic"), function=_density_minus_analytic, take_log=False, particle_type=False, display_name="density residual abs", units="g/cm**3")
                else:
                        def _particle_position_cylindrical_z_abs(field, data):
755
                                return numpy.abs(data[(PartType_Gas_to_use, "height")])
756
757
758
759
760
761
                        pf.add_field((PartType_Gas_to_use, "particle_position_cylindrical_z_abs"), function=_particle_position_cylindrical_z_abs, take_log=False, particle_type=True, units="cm")
                        # particle_type=False doesn't make sense, but is critical for PhasePlot/ProfilePlot to work for particle fields
                        # requires a change in data_objects/data_container.py: remove raise YTFieldTypeNotFound(ftype)
                        if draw_metal_map >= 1 or draw_metal_PDF == 1:
                                def _Metallicity_2(field, data):
                                        return data[(PartType_Gas_to_use, MetallicityType_to_use)]
Loic Hausammann's avatar
Gear    
Loic Hausammann committed
762
                                pf.add_field((PartType_Gas_to_use, "Metallicity_2"), function=_Metallicity_2, take_log=True, particle_type=True, display_name="Metallicity", units="")
763
                        def _Density_2_minus_analytic(field, data):
764
                                return data[(PartType_Gas_to_use, "density")] - rho_agora_disk(data[(PartType_Gas_to_use, "radius")], data[(PartType_Gas_to_use, "particle_position_cylindrical_z_abs")])
765
766
767
768
769
770
771
772
773
774
775
776
                        pf.add_field((PartType_Gas_to_use, "Density_2_minus_analytic"), function=_Density_2_minus_analytic, take_log=False, particle_type=False, display_name="Density Residual", units="g/cm**3")

                ####################################
                #      MAIN ANALYSIS ROUTINES      #
                ####################################

                # FIND CENTER AND PROJ_REGION
                v, cen = pf.h.find_max(("gas", "density")) # find the center to keep the galaxy at the center of all the images; here we assume that the gas disk is no bigger than 30 kpc in radius
                sp = pf.sphere(cen, (30.0, "kpc"))
                cen2 = sp.quantities.center_of_mass(use_gas=True, use_particles=False).in_units("kpc")
                sp2 = pf.sphere(cen2, (1.0, "kpc"))
                cen3 = sp2.quantities.max_location(("gas", "density"))
Loic Hausammann's avatar
Loic Hausammann committed
777
                center = pf.arr([cen3[1].d, cen3[2].d, cen3[3].d], 'cm') # naive usage such as YTArray([cen3[1], cen3[2], cen3[3]]) doesn't work somehow for ART-II data
778
779
                if yt_version_pre_3_2_3 == 1:
                        center = pf.arr([cen3[2].d, cen3[3].d, cen3[4].d], 'code_length') # for yt-3.2.3 or before
lhausamm's avatar
lhausamm committed
780
                if codes[code] == "GASOLINE" and dataset_num == 2 and time == 1:
781
782
                        #center = pf.arr([2.7912903206399826e+20, 1.5205303149849894e+21, 1.5398968883245956e+21], 'cm') # a temporary hack for GASOLINE (center of the most massive stellar clump)
                        center = pf.arr([ 0.09045956,  0.49277032,  0.49904659], 'code_length')
lhausamm's avatar
lhausamm committed
783
                # if codes[code] == "GADGET-3" and dataset_num == 2 and time == 1:
784
785
                #       #center = pf.arr([6.184520935812296e+21, 4.972678132728082e+21, 6.559067311284074e+21], 'cm') # a temporary hack for GADGET-3 (center of the most massive stellar clump)
                #       center = pf.arr([ 2.00426673674,  1.61153523084,  2.12564895041], 'code_length')
lhausamm's avatar
lhausamm committed
786
                if codes[code] == "ART-I" or codes[code] == "ART-II" or codes[code] == "ENZO"  or codes[code] == "RAMSES":
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
                        proj_region = pf.box(center - YTArray([figure_width, figure_width, figure_width], 'kpc'),
                                             center + YTArray([figure_width, figure_width, figure_width], 'kpc')) # projected images made using a (2*figure_width)^3 box for AMR codes
                else:
                        proj_region = pf.all_data()

                # DENSITY MAPS
                if draw_density_map == 1:
                        my_cmap = copy.copy(matplotlib.cm.get_cmap('algae'))
                        my_cmap.set_bad(my_cmap(0)) # cmap range [0, 1)
                        my_cmap.set_under(my_cmap(0))
                        for ax in range(1, 3):
                                p = ProjectionPlot(pf, ax, ("gas", "density"), center = center, data_source=proj_region, width = (figure_width, 'kpc'), weight_field = None, fontsize=9)
                                p.set_zlim(("gas", "density"), 1e-5, 1e-1)
                                p.set_cmap(("gas", "density"), my_cmap)
                                plot = p.plots[("gas", "density")]

                                plot.figure = fig_density_map[time]
                                plot.axes = grid_density_map[time][(ax-1)*len(codes)+code].axes
                                if code == 0: plot.cax = grid_density_map[time].cbar_axes[0]
                                p._setup_plots()

                        if add_nametag == 1:
                                at = AnchoredText("%s" % codes[code], loc=2, prop=dict(size=6), frameon=True)
                                grid_density_map[time][code].axes.add_artist(at)

                # TEMPERATURE MAPS
                if draw_temperature_map == 1:
                        my_cmap = copy.copy(matplotlib.cm.get_cmap('algae'))
                        my_cmap.set_bad(my_cmap(0.6)) # (log(1e4) - log(1e1)) / (log(1e6) - log(1e1)) = 0.6
                        my_cmap.set_under(my_cmap(0))
                        for ax in range(1, 3):
                                p2 = ProjectionPlot(pf, ax, ("gas", "temperature"), center = center, data_source=proj_region, width = (figure_width, 'kpc'), weight_field = ("gas", "density_squared"), fontsize=9)
                                p2.set_zlim(("gas", "temperature"), 1e1, 1e6)
                                p2.set_cmap(("gas", "temperature"), my_cmap)
                                plot2 = p2.plots[("gas", "temperature")]

                                plot2.figure = fig_temperature_map[time]
                                plot2.axes = grid_temperature_map[time][(ax-1)*len(codes)+code].axes
                                if code == 0: plot2.cax = grid_temperature_map[time].cbar_axes[0]
                                p2._setup_plots()

                        if add_nametag == 1:
                                at = AnchoredText("%s" % codes[code], loc=2, prop=dict(size=6), frameon=True)
                                grid_temperature_map[time][code].axes.add_artist(at)

                # CELL-SIZE MAPS
                if draw_cellsize_map >= 1:
                        my_cmap = copy.copy(matplotlib.cm.get_cmap('algae'))
                        my_cmap.set_bad(my_cmap(0.99999)) # 1 doesn't work!
                        my_cmap.set_under(my_cmap(0))
                        if draw_cellsize_map == 1 or draw_cellsize_map == 3:
                                for ax in range(1, 3):
                                        p25 = ProjectionPlot(pf, ax, ("index", "cell_size"), center = center, data_source=proj_region, width = (figure_width, 'kpc'), weight_field = ("index", "cell_volume_inv2"), fontsize=9)
                                        p25.set_zlim(("index", "cell_size"), 10, 1e3)
                                        p25.set_cmap(("index", "cell_size"), my_cmap)
                                        plot25 = p25.plots[("index", "cell_size")]

                                        plot25.figure = fig_cellsize_map[time]
                                        plot25.axes = grid_cellsize_map[time][(ax-1)*len(codes)+code].axes
                                        if code == 0: plot25.cax = grid_cellsize_map[time].cbar_axes[0]
                                        p25._setup_plots()

                        # Create another map with a different resolution definition if requested
                        if draw_cellsize_map == 2 or draw_cellsize_map == 3:
                                for ax in range(1, 3):
                                        if codes[code] == "ART-I" or codes[code] == "ART-II" or codes[code] == "ENZO"  or codes[code] == "RAMSES":
                                                p251 = ProjectionPlot(pf, ax, ("index", "cell_size"), center = center, data_source=proj_region, width = (figure_width, 'kpc'), \
                                                                             weight_field = ("index", "cell_volume_inv2"), fontsize=9)
                                                p251.set_zlim(("index", "cell_size"), 10, 1e3)
                                                p251.set_cmap(("index", "cell_size"), my_cmap)
                                                plot251 = p251.plots[("index", "cell_size")]
                                        else:
                                                p251 = ProjectionPlot(pf, ax, ("gas", "particle_size"), center = center, data_source=proj_region, width = (figure_width, 'kpc'), \
                                                                             weight_field = ("gas", "particle_volume_inv2"), fontsize=9)
                                                p251.set_zlim(("gas", "particle_size"), 10, 1e3)
                                                p251.set_cmap(("gas", "particle_size"), my_cmap)
                                                plot251 = p251.plots[("gas", "particle_size")]

                                        plot251.figure = fig_cellsize_map_2[time]
                                        plot251.axes = grid_cellsize_map_2[time][(ax-1)*len(codes)+code].axes
                                        if code == 0: plot251.cax = grid_cellsize_map_2[time].cbar_axes[0]
                                        p251._setup_plots()

                        if add_nametag == 1:
                                if draw_cellsize_map == 1 or draw_cellsize_map == 3:
                                        at = AnchoredText("%s" % codes[code], loc=2, prop=dict(size=6), frameon=True)
                                        grid_cellsize_map[time][code].axes.add_artist(at)
                                if draw_cellsize_map == 2 or draw_cellsize_map == 3:
                                        at = AnchoredText("%s" % codes[code], loc=2, prop=dict(size=6), frameon=True)
                                        grid_cellsize_map_2[time][code].axes.add_artist(at)

                # ELEVATION MAPS
                if draw_elevation_map == 1:
                        def _CellzElevationpc(field,data):
                                return ((data[("index", "z")] - center[2]).in_units('pc'))
                        pf.add_field(("index", "z_elevation"), function=_CellzElevationpc, units='kpc', display_name="$z$ Elevation", take_log=False)
                        my_cmap = copy.copy(matplotlib.cm.get_cmap('algae'))
                        my_cmap.set_bad(my_cmap(0.5))
                        my_cmap.set_under(my_cmap(0))
                        for ax in range(2, 3):
                                p255 = ProjectionPlot(pf, ax, ("index", "z_elevation"), center = center, data_source=proj_region, width = (figure_width, 'kpc'), weight_field = ("gas", "density"), fontsize=9)
                                p255.set_zlim(("index", "z_elevation"), -1, 1)
                                p255.set_cmap(("index", "z_elevation"), my_cmap)
                                plot255 = p255.plots[("index", "z_elevation")]

                                plot255.figure = fig_elevation_map[time]
                                plot255.axes = grid_elevation_map[time][(ax-2)*len(codes)+code].axes
                                if code == 0: plot255.cax = grid_elevation_map[time].cbar_axes[0]
                                p255._setup_plots()

                        if add_nametag == 1:
                                at = AnchoredText("%s" % codes[code], loc=2, prop=dict(size=6), frameon=True)
                                grid_elevation_map[time][code].axes.add_artist(at)

                # Z-VELOCITY MAPS
                if draw_zvel_map == 1:
                        my_cmap = copy.copy(matplotlib.cm.get_cmap('algae'))
                        my_cmap.set_bad(my_cmap(0.5))
                        my_cmap.set_under(my_cmap(0))
                        for ax in range(1, 3):
                                p26 = ProjectionPlot(pf, ax, ("gas", "velocity_z"), center = center, data_source=proj_region, width = (figure_width, 'kpc'), weight_field = ("gas", "density_squared"), fontsize=9)
                                p26.set_zlim(("gas", "velocity_z"), -1e7, 1e7)
                                p26.set_cmap(("gas", "velocity_z"), my_cmap)
                                plot26 = p26.plots[("gas", "velocity_z")]

                                plot26.figure = fig_zvel_map[time]
                                plot26.axes = grid_zvel_map[time][(ax-1)*len(codes)+code].axes
                                if code == 0: plot26.cax = grid_zvel_map[time].cbar_axes[0]
                                p26._setup_plots()

                        if add_nametag == 1:
                                at = AnchoredText("%s" % codes[code], loc=2, prop=dict(size=6), frameon=True)
                                grid_zvel_map[time][code].axes.add_artist(at)

                # METAL MAPS
                if draw_metal_map >= 1:
                        my_cmap = copy.copy(matplotlib.cm.get_cmap('algae'))
                        my_cmap.set_bad(my_cmap(0.347)) # (0.02041 - 0.01) / (0.04 - 0.01) = 0.347
                        my_cmap.set_under(my_cmap(0))
                        for ax in range(1, 3):
                                pf.field_info[("gas", "metallicity")].take_log = False
                                p26 = ProjectionPlot(pf, ax, ("gas", "metallicity"), center = center, data_source=proj_region, width = (figure_width, 'kpc'), weight_field = ("gas", "density_squared"), fontsize=9)
                                p26.set_zlim(("gas", "metallicity"), 0.01, 0.04)
                                p26.set_cmap(("gas", "metallicity"), my_cmap)
                                plot26 = p26.plots[("gas", "metallicity")]

                                plot26.figure = fig_metal_map[time]
                                plot26.axes = grid_metal_map[time][(ax-1)*len(codes)+code].axes
                                if code == 0: plot26.cax = grid_metal_map[time].cbar_axes[0]
                                p26._setup_plots()

                        if add_nametag == 1:
                                if draw_metal_map == 2:
                                        sp = pf.sphere(center, (1, "Mpc"))
                                        total_metal_mass = sp.quantities.total_quantity([("gas", "metal_mass")])
                                        at = AnchoredText("%s - %e" % (codes[code], total_metal_mass), loc=2, prop=dict(size=6), frameon=True)
                                else:
                                        at = AnchoredText("%s" % codes[code], loc=2, prop=dict(size=6), frameon=True)
                                grid_metal_map[time][code].axes.add_artist(at)

                # STELLAR MAPS
                if draw_star_map == 1 and time != 0:
                        for ax in range(1, 3):
                                p27 = ParticleProjectionPlot(pf, ax, (PartType_Star_to_use, "particle_mass"), center = center, data_source=proj_region, width = (figure_width, 'kpc'), weight_field = None, fontsize=9)
                                p27.set_unit((PartType_Star_to_use, "particle_mass"), 'Msun')
                                p27.set_zlim((PartType_Star_to_use, "particle_mass"), 1e4, 1e7)
                                p27.set_cmap((PartType_Star_to_use, "particle_mass"), 'algae')
                                p27.set_buff_size(400) # default is 800
                                p27.set_colorbar_label((PartType_Star_to_use, "particle_mass"), "Stellar Mass Per Pixel ($\mathrm{M}_{\odot}$)")
                                plot27 = p27.plots[(PartType_Star_to_use, "particle_mass")]
                                # p27 = ParticleProjectionPlot(pf, ax, (PartType_Star_to_use, "particle_velocity_cylindrical_theta"), center = center, data_source=proj_region, width = (figure_width, 'kpc'), weight_field = (PartType_Star_to_use, "particle_mass"), fontsize=9)
                                # p27.set_unit((PartType_Star_to_use, "particle_velocity_cylindrical_theta"), 'km/s')
                                # p27.set_buff_size(400)
                                # p27.set_colorbar_label((PartType_Star_to_use, "particle_velocity_cylindrical_theta"), "Rotational Velocity (km/s)")
                                # plot27 = p27.plots[(PartType_Star_to_use, "particle_velocity_cylindrical_theta")]

                                plot27.figure = fig_star_map[time]
                                plot27.axes = grid_star_map[time][(ax-1)*len(codes)+code].axes
                                if code == 0: plot27.cax = grid_star_map[time].cbar_axes[0]
                                p27._setup_plots()

                        if add_nametag == 1:
                                at = AnchoredText("%s" % codes[code], loc=2, prop=dict(size=6), frameon=True)
                                grid_star_map[time][code].axes.add_artist(at)

                # STELLAR CLUMP STATISTICS AND ANNOTATED STELLAR MAPS
                if draw_star_clump_stats >= 1 and time != 0:
                        # For make yt's HaloCatalog to work with non-cosmological dataset, a fix needed to be applied to analysis_modules/halo_finding/halo_objects.py: self.period = ds.arr([3.254, 3.254, 3.254], 'Mpc')
                        pf.hubble_constant = 0.71; pf.omega_lambda = 0.73; pf.omega_matter = 0.27; pf.omega_curvature = 0.0; pf.current_redshift = 0.0 # another trick to make HaloCatalog work especially for ART-I dataset
                        if os.path.exists("./halo_catalogs/hop_%s_%05d/hop_%s_%05d.0.h5" % (codes[code], times[time], codes[code], times[time])) == False:
                                hc = HaloCatalog(data_ds=pf, finder_method='hop', output_dir="./halo_catalogs/hop_%s_%05d" % (codes[code], times[time]), \
                                                         finder_kwargs={'threshold': 2e8, 'dm_only': False, 'ptype': PartType_Star_to_use})
#                                                        finder_kwargs={'threshold': 2e5, 'dm_only': False, 'ptype': "all"})
                                hc.add_filter('quantity_value', 'particle_mass', '>', 2.6e6, 'Msun') # exclude halos with less than 30 particles
                                hc.add_filter('quantity_value', 'particle_mass', '<', 2.6e8, 'Msun') # exclude the most massive halo (threshold 1e8.4 is hand-picked, so one needs to be careful!)
                                hc.create()
                        halo_ds = load("./halo_catalogs/hop_%s_%05d/hop_%s_%05d.0.h5" % (codes[code], times[time], codes[code], times[time]))
                        hc = HaloCatalog(halos_ds=halo_ds, output_dir="./halo_catalogs/hop_%s_%05d" % (codes[code], times[time]))
                        hc.load()

                        halo_ad = hc.halos_ds.all_data()
                        star_clump_masses_hop[time].append(np.log10(halo_ad['particle_mass'][:].in_units("Msun")))

                        if os.path.exists("./halo_catalogs/fof_%s_%05d/fof_%s_%05d.0.h5" % (codes[code], times[time], codes[code], times[time])) == False:
                                hc2 = HaloCatalog(data_ds=pf, finder_method='fof', output_dir="./halo_catalogs/fof_%s_%05d" % (codes[code], times[time]), \
                                                          finder_kwargs={'link': 0.0025, 'dm_only': False, 'ptype': PartType_Star_to_use})
#                                                         finder_kwargs={'link': 0.02, 'dm_only': False, 'ptype': "all"})
                                hc2.add_filter('quantity_value', 'particle_mass', '>', 2.6e6, 'Msun') # exclude halos with less than 30 particles
                                hc2.add_filter('quantity_value', 'particle_mass', '<', 2.6e8, 'Msun') # exclude the most massive halo (threshold 1e8.4 is hand-picked, so one needs to be careful!)
                                hc2.create()
                        halo_ds2 = load("./halo_catalogs/fof_%s_%05d/fof_%s_%05d.0.h5" % (codes[code], times[time], codes[code], times[time]))
                        hc2 = HaloCatalog(halos_ds=halo_ds2, output_dir="./halo_catalogs/fof_%s_%05d" % (codes[code], times[time]))
                        hc2.load()

For faster browsing, not all history is shown. View entire blame