RAMSES_compound_data.hh

/*
	************************************************************************
	This file is part of RAMSEStools++
	a suite of C++ classes and routines to postprocess data
	generated by the simulation code RAMSES by R. Teyssier

    Copyright (C) 2009  Oliver Hahn, ojha@gmx.de

	************************************************************************
*/

#ifndef __RAMSES_COMPOUND_DATA_HH
#define __RAMSES_COMPOUND_DATA_HH

#include "RAMSES_amr_data.hh"
#include "RAMSES_hydro_data.hh"
#include "RAMSES_particle_data.hh"
#include "vtkMultiprocessController.h"
#include "lattice_RAMSES.hh"

#define R_CHECK_BIT(var,pos) ((var) & (1<<(pos)))

namespace RAMSES{



inline void mpi_distribute_domains( int ndomains, std::vector<int>& mycpus, bool silent=false )
{
  int rank=0;
  int size=1;
  //... determine meta-domain-decomposition ...//
  if(vtkMultiprocessController::GetGlobalController()!=NULL){
     rank = vtkMultiProcessController::GetGlobalController()->GetLocalProcessId();
     size= vtkMultiProcessController::GetGlobalController()->GetNumberOfProcesses();
  }    

  //... determine meta-domain-decomposition ...//
  int npp = (int)((float)ndomains/size);
  
	if( rank == size-1 )
    npp = ndomains-(rank*npp);
	
	if(!silent)
		std::cout << "* Task " << rank << ": working on domains "
      << 1+npp*rank << " - " << npp*(rank+1) << std::endl;
  
	for( int i=1+npp*rank; i<=npp*(rank+1); ++i )
      mycpus.push_back(i);
  if(vtkMultiprocessController::GetGlobalController()!=NULL){
    vtkMultiProcessController::GetGlobalController()->Barrier();
  }
}
	
class compound_data{
public:

	typedef lattice<int>::tree tree_t;

	typedef PART::multi_domain_data< tree_t, double > mddouble; 
	typedef PART::multi_domain_data< tree_t, unsigned > mdunsigned; 
	typedef PART::multi_domain_data< tree_t, int > mdint;
		
	typedef HYDRO::multi_domain_data< tree_t, HYDRO::data<tree_t,double>, double > mdhydro;
	typedef HYDRO::multi_domain_data< tree_t, HYDRO::empty_data<tree_t,double>, double > mddummy;

	enum type{
		particle = 1L << 0,
		grid     = 1L << 1
	};

protected:

class tree_wrap
	{
		public:
			std::vector< tree_t* > m_trees;
			unsigned m_ntrees;
		
		tree_wrap( RAMSES::snapshot *psnap, const std::vector<int>& mycpus )
		: m_ntrees(0)
		{
			for( unsigned i=0; i<mycpus.size(); ++i )
			{
				m_trees.push_back( new RAMSES::lattice<int>::tree( *psnap, mycpus[i], psnap->m_header.levelmax ) );//maxlevel , minlevel ) );
				m_trees.back()->read();
				++m_ntrees;
			}
		}
		
		~tree_wrap()
		{
			for( unsigned i=0; i<m_ntrees; ++i )
				delete m_trees[i];
		}
		
		
		std::vector< tree_t* >& operator*( void ){ return m_trees; }
	};
	
public:
	
	class domain_distributor
	{
		public:
			std::vector<int> m_cpus;
		
			domain_distributor( RAMSES::snapshot *psnap, bool silent=false )
			{
				RAMSES::mpi_distribute_domains( psnap->m_header.ncpu, m_cpus, silent );
				//RAMSES::mpi_distribute_domains( 4, m_cpus, silent );
			}
			
			domain_distributor( const domain_distributor& d )
			: m_cpus( d.m_cpus )
			{ }
			
			int operator[]( unsigned i ) { return m_cpus[i]; }
			unsigned size( void ) { return m_cpus.size(); }
	};

public:
	std::string m_snapname;
	RAMSES::snapshot *m_prsnap;

	domain_distributor m_mycpus;

	tree_wrap m_ptrees;

	 mddouble
		m_part_pos_x, 
		m_part_pos_y, 
		m_part_pos_z,
		m_part_vel_x, 
		m_part_vel_y, 
		m_part_vel_z,
		m_part_age,
		m_part_mass;

	mdunsigned
		m_part_ids;

	mdint
		m_part_reflevel;

	mdhydro
		m_gas_rho,
		m_gas_vel_x,
		m_gas_vel_y,
		m_gas_vel_z,
		m_gas_pressure,
		m_gas_metals;
		
	std::vector< unsigned > m_nparticles;
	
	int m_type;
	
	std::vector<double>
		m_pbbox_xmin,
		m_pbbox_xmax,
		m_pbbox_ymin,
		m_pbbox_ymax,
		m_pbbox_zmin,
		m_pbbox_zmax;
	
	bool 
		m_bhave_stars,
		m_bhave_metals;


public:
	compound_data( RAMSES::snapshot &psnap, int type )
	:   m_prsnap( &psnap ),
		m_mycpus( m_prsnap/*mycpus*/ ), m_ptrees( m_prsnap, m_mycpus.m_cpus ),
		m_part_pos_x   ( *m_prsnap, *m_ptrees ),
		m_part_pos_y   ( *m_prsnap, *m_ptrees ), 
		m_part_pos_z   ( *m_prsnap, *m_ptrees ),
		m_part_vel_x   ( *m_prsnap, *m_ptrees ), 
		m_part_vel_y   ( *m_prsnap, *m_ptrees ), 
		m_part_vel_z   ( *m_prsnap, *m_ptrees ),
		m_part_age     ( *m_prsnap, *m_ptrees ),
		m_part_mass    ( *m_prsnap, *m_ptrees ),
		m_part_ids     ( *m_prsnap, *m_ptrees ),
		m_part_reflevel( *m_prsnap, *m_ptrees ),
		m_gas_rho      ( *m_prsnap, *m_ptrees ),
		m_gas_vel_x    ( *m_prsnap, *m_ptrees ),
		m_gas_vel_y    ( *m_prsnap, *m_ptrees ),
		m_gas_vel_z    ( *m_prsnap, *m_ptrees ),
		m_gas_pressure ( *m_prsnap, *m_ptrees ),
		m_gas_metals	 ( *m_prsnap, *m_ptrees ),
		m_type( type )
	{
		if( R_CHECK_BIT( m_type, 0 ) )
		{
			if(vtkMultiProcessController::GetGlobalController()==NULL || vtkMultiProcessController::GetGlobalController()->GetLocalProcessId()==0 )
				std::cout << "* Reading particle data..." << std::endl;
			
			m_part_pos_x.get_var("position_x");
			m_part_pos_y.get_var("position_y");
			m_part_pos_z.get_var("position_z");
		
			m_part_vel_x.get_var("velocity_x");
			m_part_vel_y.get_var("velocity_y");
			m_part_vel_z.get_var("velocity_z");
		
			if(! m_part_age.get_var("age") )
				m_bhave_stars = false;
			else {
				m_bhave_stars = true;
			}

			
			m_part_ids.get_var("particle_ID");
			m_part_mass.get_var("mass");
			m_part_reflevel.get_var("refinement_level");
			
			for( unsigned i=0; i<m_mycpus.size(); ++i )
				m_nparticles.push_back( m_part_reflevel.size(i) );
			
			
			m_pbbox_xmin.assign( m_mycpus.size(), +1e30 );
			m_pbbox_xmax.assign( m_mycpus.size(), -1e30 );
			m_pbbox_ymin.assign( m_mycpus.size(), +1e30 );
			m_pbbox_ymax.assign( m_mycpus.size(), -1e30 );
			m_pbbox_zmin.assign( m_mycpus.size(), +1e30 );
			m_pbbox_zmax.assign( m_mycpus.size(), -1e30 );
			
			//... set all particles to the finest level ...//
			for( unsigned i=0; i<m_mycpus.size(); ++i )
				for( unsigned ip = 0; ip < m_nparticles[i]; ++ip ){
					m_part_reflevel(i,ip) = m_prsnap->m_header.levelmax;
					
					if( m_part_pos_x(i,ip) > m_pbbox_xmax[i] )
						m_pbbox_xmax[i] = m_part_pos_x(i,ip);
					if( m_part_pos_x(i,ip) < m_pbbox_xmin[i] )
						m_pbbox_xmin[i] = m_part_pos_x(i,ip);
					
					if( m_part_pos_y(i,ip) > m_pbbox_ymax[i] )
						m_pbbox_ymax[i] = m_part_pos_y(i,ip);
					if( m_part_pos_y(i,ip) < m_pbbox_ymin[i] )
						m_pbbox_ymin[i] = m_part_pos_y(i,ip);
					
					if( m_part_pos_z(i,ip) > m_pbbox_zmax[i] )
						m_pbbox_zmax[i] = m_part_pos_z(i,ip);
					if( m_part_pos_z(i,ip) < m_pbbox_zmin[i] )
						m_pbbox_zmin[i] = m_part_pos_z(i,ip);
				}
		}else{
			for( unsigned i=0; i<m_mycpus.size(); ++i )
				m_nparticles.push_back(0);
		}
		
		if( R_CHECK_BIT( m_type, 1 ) )
		{
			if(vtkMultiProcessController::GetGlobalController()==NULL ||vtkMultiProcessController::GetGlobalController()->GetLocalProcessId()==0 )
				std::cout << "* Reading mesh data..." << std::endl;
			
			m_gas_rho.get_var("density");
			m_gas_vel_x.get_var("velocity_x");
			m_gas_vel_y.get_var("velocity_y");
			m_gas_vel_z.get_var("velocity_z");
			m_gas_pressure.get_var("pressure");
			if( !m_gas_metals.get_var("metallicity") )
				m_bhave_metals = false;
			else
				m_bhave_metals = true;

		}
    if(vtkMultiProcessController::GetGlobalController()!=NULL){
      vtkMultiProcessController::GetGlobalController()->Barrier();
    }

		if(vtkMultiProcessController::GetGlobalController()==NULL || vtkMultiProcessController::GetGlobalController()->GetLocalProcessId()==0 )
			std::cout << "* Finished reading data." << std::endl;
	}
	
	#define R_SQR(x) ((x)*(x))
	bool particles_intersect_ball( unsigned idomain, double xc, double yc, double zc, double r )
	{
	double dmin = 0;
		
		//.. x ..//		
		if( xc < m_pbbox_xmin[idomain] ) dmin += R_SQR( xc - m_pbbox_xmin[idomain] ); else
		if( xc > m_pbbox_xmax[idomain] ) dmin += R_SQR( xc - m_pbbox_xmax[idomain] );
			
		//.. y ..//
		if( yc < m_pbbox_ymin[idomain] ) dmin += R_SQR( yc - m_pbbox_ymin[idomain] ); else
		if( yc > m_pbbox_ymax[idomain] ) dmin += R_SQR( yc - m_pbbox_ymax[idomain] );
		
		//.. x ..//
		if( zc < m_pbbox_zmin[idomain] ) dmin += R_SQR( zc - m_pbbox_zmin[idomain] ); else
		if( zc > m_pbbox_zmax[idomain] ) dmin += R_SQR( zc - m_pbbox_zmax[idomain] );
		
		if( dmin <= r*r ) return true;
		return false;
	}
	#undef R_SQR
	
	
	mddummy* create_particle_densities( lattice< int >& domain_level, unsigned search_level, int ptype = PART::ptype_star  )
	{
		if( !R_CHECK_BIT( m_type, 0 ) )
			throw std::runtime_error( "attempted a call to RAMSES::compound_data::create_particle_densities, but no particle data was loaded!" );
		
		/*if( R_CHECK_BIT( ptype, RAMSES::STAR_BIT ) && !m_bhave_stars )
		{
			std::cerr	<< "Warning: Stellar density was requested, but no star particles are present.\n"
						<< "         Disabling star particles in density field...\n";
			ptype -= PART::ptype_star;
		}*/

		
		mddummy* rho_part = new mddummy( *m_prsnap, *m_ptrees );
		lattice< int >::tree::iterator grid_it;

		
		//--------------------------------------------------------------//
		//--- assign baryonic particles to the mesh using simple NGP ---//
		//--------------------------------------------------------------//
		
		char cell_ind(0);
		double dx = 1./pow(2.0,search_level+1);
		double dV = dx*dx*dx;
		unsigned icell_cpu(0);
		
		
		for( unsigned i=0; i<m_mycpus.size(); ++i ){
			for( unsigned ip = 0; ip < m_nparticles[i]; ++ip ){
				if( PART::is_of_type( m_part_age(i,ip), m_part_ids(i,ip), ptype ) ){
					//&& part_reflevel(i,ip) == search_level+1 ){
					bool ifnd = domain_level.find_cell( m_part_pos_x(i,ip), m_part_pos_y(i,ip), m_part_pos_z(i,ip), grid_it, cell_ind, icell_cpu );
					if( ifnd ){
						(*rho_part)(icell_cpu, grid_it, cell_ind) += m_part_mass(i,ip)/dV;
						
						double ddV = dV;
						
						//... and add to all finer grids associated below this one
						while( grid_it.get_level() <= m_prsnap->m_header.levelmax )
						{
							ddV = ddV/8.0;
							char cidx = grid_it.get_child_from_pos( m_part_pos_x(i,ip), m_part_pos_y(i,ip), m_part_pos_z(i,ip) );
							
							if( !grid_it.is_refined(cidx) )
								break;
							
							grid_it.to_child(cidx);
							(*rho_part)(icell_cpu, grid_it, cidx) += m_part_mass(i,ip)/ddV;
						}
					}
				}
			}
		}
		
		return rho_part;
	}
	
};
}// namespace RAMSES

#undef R_CHECK_BIT


#endif