2022-04-13 13:20:32 +02:00
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# import vtk
import sys
import time
from pathlib import Path
from sys import getsizeof
import h5py
import numba
import numpy as np
import scipy . spatial
@numba.njit ( )
def determine_desired_range ( offset , minimum , upper_limit_bottom , lower_limit_top , maximum ) :
a = minimum
b = maximum
if offset < 0 :
a = lower_limit_top
elif offset > 0 :
b = upper_limit_bottom
return a , b
@numba.njit ( )
def find_coordinates_to_move ( minimum , maximum , ratio , x_offset , y_offset , z_offset , move_candidates ) :
coordinates_to_move = [ ]
x_start , x_end = determine_desired_range ( x_offset , minimum , upper_limit_bottom , lower_limit_top , maximum )
y_start , y_end = determine_desired_range ( y_offset , minimum , upper_limit_bottom , lower_limit_top , maximum )
z_start , z_end = determine_desired_range ( z_offset , minimum , upper_limit_bottom , lower_limit_top , maximum )
for particle in move_candidates :
point = particle [ 0 : 3 ]
if x_start < = point [ 0 ] < = x_end and y_start < = point [ 1 ] < = y_end and z_start < = point [ 2 ] < = z_end :
coordinates_to_move . append ( particle )
return coordinates_to_move
2022-05-23 15:30:43 +02:00
# directory = Path(r"/home/ben/sims/swiftsim/examples/zoom_tests/auriga6_halo7_8_10/")
directory = Path ( r " /home/ben/sims/data_swift/monofonic_tests/DB2_128_100/ " )
2022-04-13 13:20:32 +02:00
# file = h5py.File(directory / "auriga6_halo7_8_9.hdf5", "r")
2022-05-23 15:30:43 +02:00
# for filename in sorted(directory.glob("output_*.hdf5")):
for filename in sorted ( directory . glob ( " output_0004.hdf5 " ) ) :
2022-04-13 13:20:32 +02:00
print ( filename )
file = h5py . File ( str ( filename ) , " r " )
Header = file [ ' Header ' ]
highres_coordinates = file [ " PartType1 " ] [ " Coordinates " ] [ : ] # for cdm particles
highres_names = file [ " PartType1 " ] [ " ParticleIDs " ] [ : ]
highres_velocities = file [ " PartType1 " ] [ " Velocities " ] [ : ]
highres_masses = file [ ' PartType1 ' ] [ ' Masses ' ] [ : ]
2022-04-25 15:09:20 +02:00
highres_group_ids = file [ ' PartType1 ' ] [ ' FOFGroupIDs ' ] [ : ]
2022-04-13 13:20:32 +02:00
highres_absolute_velo = np . sqrt ( np . sum ( highres_velocities * * 2 , axis = 1 ) )
lowres_coordinates = file [ " PartType2 " ] [ " Coordinates " ] [ : ] # for cdm particles
lowres_names = file [ " PartType2 " ] [ " ParticleIDs " ] [ : ]
lowres_velocities = file [ " PartType2 " ] [ " Velocities " ] [ : ]
lowres_masses = file [ ' PartType2 ' ] [ ' Masses ' ] [ : ]
2022-04-25 15:09:20 +02:00
lowres_group_ids = file [ ' PartType2 ' ] [ ' FOFGroupIDs ' ] [ : ]
2022-04-13 13:20:32 +02:00
lowres_absolute_velo = np . sqrt ( np . sum ( lowres_velocities * * 2 , axis = 1 ) )
original_coordinates = np . concatenate ( ( highres_coordinates , lowres_coordinates ) )
# if "auriga" in str(filename):
# original_coordinates /= 1000
2022-04-25 15:09:20 +02:00
# print(original_coordinates.mean())
# print(original_coordinates.min())
# print(original_coordinates.max())
2022-04-13 13:20:32 +02:00
# print(file['Header'])
# print(list(file['Units'].attrs.items()))
# exit()
names = np . concatenate ( ( highres_names , lowres_names ) )
velocities = np . concatenate ( ( highres_velocities , lowres_velocities ) )
masses = np . concatenate ( ( highres_masses , lowres_masses ) )
2022-04-25 15:09:20 +02:00
group_ids = np . concatenate ( ( highres_group_ids , lowres_group_ids ) )
2022-04-13 13:20:32 +02:00
absolute_velo = np . concatenate ( ( highres_absolute_velo , lowres_absolute_velo ) )
# for bla in [original_coordinates, names, velocities, masses, absolute_velo]:
# print(bla.shape)
original_data = np . vstack ( [
original_coordinates [ : : , 0 ] ,
original_coordinates [ : : , 1 ] ,
original_coordinates [ : : , 2 ] ,
names ,
velocities [ : : , 0 ] ,
velocities [ : : , 1 ] ,
velocities [ : : , 2 ] ,
masses ,
2022-04-25 15:09:20 +02:00
group_ids ,
2022-04-13 13:20:32 +02:00
absolute_velo ,
] ) . T
print ( original_data . shape )
assert ( original_coordinates == original_data [ : : , 0 : 3 ] ) . all ( )
boundaries = Header . attrs [ ' BoxSize ' ] # BoxLength for e5 boxes depends on Nres, 2.36438 for 256, 4.72876 for 512.
2022-04-25 15:09:20 +02:00
print ( boundaries , len ( highres_names ) )
2022-04-13 13:20:32 +02:00
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 ( )
def find_move_candidates ( ) :
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
move_candidates = find_move_candidates ( )
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)
2022-04-25 15:09:20 +02:00
# print(densities)
# print(densities.shape)
2022-04-13 13:20:32 +02:00
all_data = np . column_stack ( [ list ( range ( densities . shape [ 0 ] ) ) , transformed_data , densities , alt_densities ] )
2022-04-25 15:09:20 +02:00
# print(all_data.shape)
# print(original_data.shape[0])
2022-04-13 13:20:32 +02:00
export_data = all_data [ : original_data . shape [ 0 ] ]
2022-04-25 15:09:20 +02:00
# print(export_data.shape)
2022-04-13 13:20:32 +02:00
# 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 ( directory / f " out_ { filename . with_suffix ( ' .csv ' ) . name } " ,
export_data ,
delimiter = " , " ,
fmt = " %.3f " ,
2022-04-25 15:09:20 +02:00
header = " num,x,y,z,name,vx,vy,vz,masse,groupid,v,density,density_alt " )
2022-04-13 13:20:32 +02:00
file . close ( )