1
0
Fork 0
mirror of https://github.com/glatterf42/sims_python_files.git synced 2024-09-19 16:13:45 +02:00

extract methods

this commit shouldn't change any actual code apart from reordering it
and extracting parts of it into methods
This commit is contained in:
Lukas Winkler 2022-05-23 14:59:34 +02:00
parent 05e6be9921
commit 5da3f1ad26
Signed by: lukas
GPG key ID: 54DE4D798D244853

View file

@ -41,21 +41,34 @@ def find_coordinates_to_move(minimum, maximum, ratio, x_offset, y_offset, z_offs
return coordinates_to_move return coordinates_to_move
# file = h5py.File(directory / "auriga6_halo7_8_9.hdf5", "r") @numba.njit()
def find_move_candidates(original_data, minimum, maximum, lower_limit_top, upper_limit_bottom):
move_candidates = []
print("finding move candidates")
for particle in original_data:
point = particle[0:3]
if (
minimum <= point[0] <= upper_limit_bottom or
lower_limit_top <= point[0] <= maximum or
minimum <= point[1] <= upper_limit_bottom or
lower_limit_top <= point[1] <= maximum or
minimum <= point[2] <= upper_limit_bottom or
lower_limit_top <= point[2] <= maximum
):
move_candidates.append(particle)
# print(point)
return move_candidates
for filename in sys.argv[1:]:
filename = Path(filename) def read_file(filename):
print(filename)
file = h5py.File(str(filename), "r") file = h5py.File(str(filename), "r")
Header = file['Header'] Header = file['Header']
highres_coordinates = file["PartType1"]["Coordinates"][:] # for cdm particles highres_coordinates = file["PartType1"]["Coordinates"][:] # for cdm particles
highres_names = file["PartType1"]["ParticleIDs"][:] highres_names = file["PartType1"]["ParticleIDs"][:]
highres_velocities = file["PartType1"]["Velocities"][:] highres_velocities = file["PartType1"]["Velocities"][:]
highres_masses = file['PartType1']['Masses'][:] highres_masses = file['PartType1']['Masses'][:]
highres_group_ids = file['PartType1']['FOFGroupIDs'][:] highres_group_ids = file['PartType1']['FOFGroupIDs'][:]
highres_absolute_velo = np.sqrt(np.sum(highres_velocities ** 2, axis=1)) highres_absolute_velo = np.sqrt(np.sum(highres_velocities ** 2, axis=1))
if "PartType2" in file: if "PartType2" in file:
lowres_coordinates = file["PartType2"]["Coordinates"][:] # for cdm particles lowres_coordinates = file["PartType2"]["Coordinates"][:] # for cdm particles
lowres_names = file["PartType2"]["ParticleIDs"][:] lowres_names = file["PartType2"]["ParticleIDs"][:]
@ -77,7 +90,7 @@ for filename in sys.argv[1:]:
masses = highres_masses masses = highres_masses
group_ids = highres_group_ids group_ids = highres_group_ids
absolute_velo = highres_absolute_velo absolute_velo = highres_absolute_velo
file.close()
# if "auriga" in str(filename): # if "auriga" in str(filename):
# original_coordinates /= 1000 # original_coordinates /= 1000
# print(original_coordinates.mean()) # print(original_coordinates.mean())
@ -103,124 +116,117 @@ for filename in sys.argv[1:]:
]).T ]).T
print(original_data.shape) print(original_data.shape)
assert (original_coordinates == original_data[::, 0:3]).all() assert (original_coordinates == original_data[::, 0:3]).all()
return Header, highres_names, original_data
boundaries = Header.attrs['BoxSize'] # BoxLength for e5 boxes depends on Nres, 2.36438 for 256, 4.72876 for 512.
print(boundaries, len(highres_names))
if not boundaries.shape:
boundaries = np.array([boundaries] * 3)
offsets = [-1, 0, 1]
transformed_data = original_data[:]
number_of_time_that_points_have_been_found = 0
# assumes cube form and 0.1 as desired ratio to move
minimum = 0.0
maximum = max(boundaries)
ratio = 0.1
box_length = maximum - minimum
range_to_move = 0.1 * box_length
upper_limit_bottom = minimum + range_to_move
lower_limit_top = maximum - range_to_move
print("Find candidates to move...")
@numba.njit() # file = h5py.File(directory / "auriga6_halo7_8_9.hdf5", "r")
def find_move_candidates(): def main():
move_candidates = [] for filename in sys.argv[1:]:
print("finding move candidates") filename = Path(filename)
for particle in original_data: print(filename)
point = particle[0:3] Header, highres_names, original_data = read_file(filename)
if (
minimum <= point[0] <= upper_limit_bottom or boundaries = Header.attrs[
lower_limit_top <= point[0] <= maximum or 'BoxSize'] # BoxLength for e5 boxes depends on Nres, 2.36438 for 256, 4.72876 for 512.
minimum <= point[1] <= upper_limit_bottom or print(boundaries, len(highres_names))
lower_limit_top <= point[1] <= maximum or if not boundaries.shape:
minimum <= point[2] <= upper_limit_bottom or boundaries = np.array([boundaries] * 3)
lower_limit_top <= point[2] <= maximum offsets = [-1, 0, 1]
): transformed_data = original_data[:]
move_candidates.append(particle) number_of_time_that_points_have_been_found = 0
# print(point)
return move_candidates # assumes cube form and 0.1 as desired ratio to move
minimum = 0.0
maximum = max(boundaries)
ratio = 0.1
box_length = maximum - minimum
range_to_move = 0.1 * box_length
upper_limit_bottom = minimum + range_to_move
lower_limit_top = maximum - range_to_move
print("Find candidates to move...")
move_candidates = find_move_candidates(original_data, minimum, maximum, lower_limit_top, upper_limit_bottom)
move_candidates = np.array(move_candidates)
print("...done.")
for x in offsets:
for y in offsets:
for z in offsets:
if (x, y, z) == (0, 0, 0):
continue
moved_coordinates = find_coordinates_to_move(minimum, maximum, ratio, x, y, z, move_candidates)
# print(moved_coordinates)
moved_coordinates = np.array(moved_coordinates)
# if not moved_coordinates:
# print(f"nothing moved in {(x,y,z)}")
# continue
moved_coordinates[::, 0] += x * boundaries[0]
moved_coordinates[::, 1] += y * boundaries[1]
moved_coordinates[::, 2] += z * boundaries[2]
transformed_data = np.vstack((transformed_data, moved_coordinates))
number_of_time_that_points_have_been_found += 1
print(f"Points found: {number_of_time_that_points_have_been_found}/26...")
# assert coordinates.shape[0] == original_coordinates.shape[0] * 3 ** 3 #check that the new space has the shape we want it to have
num_nearest_neighbors = 40
print("Building 3d-Tree for all particles...")
coordinates = transformed_data[::, 0:3]
print(coordinates.shape)
tree = scipy.spatial.KDTree(coordinates)
print(getsizeof(tree) / 1024, "KB")
print("...done.")
print("Searching neighbours...")
a = time.perf_counter_ns()
distances, indices = tree.query([coordinates], k=num_nearest_neighbors, workers=6)
# shape of closest_neighbours: (1, xxxx, 40)
b = time.perf_counter_ns()
print("...found neighbours.")
print(f"took {(b - a) / 1000 / 1000:.2f} ms")
distances = distances[0] # to (xxxx, 40)
indices = indices[0] # to (xxxx, 40)
print(distances.shape)
print(indices.shape)
print(indices)
mass_array = []
print("fetching masses")
for subindices in indices: # subindices is (40)
# can maybe be optimized to remove loop
masses = transformed_data[subindices, 7]
mass_array.append(masses)
mass_array = np.array(mass_array)
print("finished fetching masses")
# print(closest_neighbours, indices)
# print(indices)
# densities = num_nearest_neighbors * mass_per_particle / np.mean(closest_neighbours, axis=1) ** 3
total_masses = np.sum(mass_array, axis=1)
densities = total_masses / np.mean(distances, axis=1) ** 3
alt_densities = total_masses / np.max(distances, axis=1) ** 3
# print(closest_neighbours.shape)
# print(densities)
# print(densities.shape)
all_data = np.column_stack([list(range(densities.shape[0])), transformed_data, densities, alt_densities])
# print(all_data.shape)
# print(original_data.shape[0])
export_data = all_data[:original_data.shape[0]]
# print(export_data.shape)
# all_data = np.append(coordinates, velocities, axis=1)
# all_data = np.column_stack((all_data, absolute_velo, names, densities))
# sorted_index = np.argsort(all_data[::, 7], kind="stable")
# all_data = all_data[sorted_index, :]
# np.savetxt("out_"+filename.with_suffix(".csv").name, all_data[indices], delimiter=",", fmt="%.3f", header="x,y,z,vx,vy,vz,v,name") #if indices are needed
np.savetxt(f"out_{filename.with_suffix('.csv').name}",
export_data,
delimiter=",",
fmt="%.3f",
header="num,x,y,z,name,vx,vy,vz,masse,groupid,v,density,density_alt")
move_candidates = find_move_candidates() if __name__ == '__main__':
move_candidates = np.array(move_candidates) main()
print("...done.")
for x in offsets:
for y in offsets:
for z in offsets:
if (x, y, z) == (0, 0, 0):
continue
moved_coordinates = find_coordinates_to_move(minimum, maximum, ratio, x, y, z, move_candidates)
# print(moved_coordinates)
moved_coordinates = np.array(moved_coordinates)
# if not moved_coordinates:
# print(f"nothing moved in {(x,y,z)}")
# continue
moved_coordinates[::, 0] += x * boundaries[0]
moved_coordinates[::, 1] += y * boundaries[1]
moved_coordinates[::, 2] += z * boundaries[2]
transformed_data = np.vstack((transformed_data, moved_coordinates))
number_of_time_that_points_have_been_found += 1
print(f"Points found: {number_of_time_that_points_have_been_found}/26...")
# assert coordinates.shape[0] == original_coordinates.shape[0] * 3 ** 3 #check that the new space has the shape we want it to have
num_nearest_neighbors = 40
print("Building 3d-Tree for all particles...")
coordinates = transformed_data[::, 0:3]
print(coordinates.shape)
tree = scipy.spatial.KDTree(coordinates)
print(getsizeof(tree) / 1024, "KB")
print("...done.")
print("Searching neighbours...")
a = time.perf_counter_ns()
distances, indices = tree.query([coordinates], k=num_nearest_neighbors, workers=6)
# shape of closest_neighbours: (1, xxxx, 40)
b = time.perf_counter_ns()
print("...found neighbours.")
print(f"took {(b - a) / 1000 / 1000:.2f} ms")
distances = distances[0] # to (xxxx, 40)
indices = indices[0] # to (xxxx, 40)
print(distances.shape)
print(indices.shape)
print(indices)
mass_array = []
print("fetching masses")
for subindices in indices: # subindices is (40)
# can maybe be optimized to remove loop
masses = transformed_data[subindices, 7]
mass_array.append(masses)
mass_array = np.array(mass_array)
print("finished fetching masses")
# print(closest_neighbours, indices)
# print(indices)
# densities = num_nearest_neighbors * mass_per_particle / np.mean(closest_neighbours, axis=1) ** 3
total_masses = np.sum(mass_array, axis=1)
densities = total_masses / np.mean(distances, axis=1) ** 3
alt_densities = total_masses / np.max(distances, axis=1) ** 3
# print(closest_neighbours.shape)
# print(densities)
# print(densities.shape)
all_data = np.column_stack([list(range(densities.shape[0])), transformed_data, densities, alt_densities])
# print(all_data.shape)
# print(original_data.shape[0])
export_data = all_data[:original_data.shape[0]]
# print(export_data.shape)
# all_data = np.append(coordinates, velocities, axis=1)
# all_data = np.column_stack((all_data, absolute_velo, names, densities))
# sorted_index = np.argsort(all_data[::, 7], kind="stable")
# all_data = all_data[sorted_index, :]
# np.savetxt("out_"+filename.with_suffix(".csv").name, all_data[indices], delimiter=",", fmt="%.3f", header="x,y,z,vx,vy,vz,v,name") #if indices are needed
np.savetxt(f"out_{filename.with_suffix('.csv').name}",
export_data,
delimiter=",",
fmt="%.3f",
header="num,x,y,z,name,vx,vy,vz,masse,groupid,v,density,density_alt")
file.close()