first commit

.gitignore

@@ -0,0 +1,5 @@
+.DS_Store
+batchfile.slurm
+crosstest_batchfile.slurm
+tunning_batchfile.slurm
+__pycache__/

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2022 Pablo Villanueva Domingo
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

PS_files/k_values.txt

@@ -0,0 +1,79 @@
+3.559856489878782115e-01
+6.031648866026231293e-01
+8.579160745827347778e-01
+1.113464310595710716e+00
+1.374039914321364408e+00
+1.619154704570298309e+00
+1.870686673430271840e+00
+2.123290040706380832e+00
+2.378144976376236119e+00
+2.630489925835516996e+00
+2.878653162679494670e+00
+3.131144366275339319e+00
+3.388159578333816313e+00
+3.638741848210948149e+00
+3.888315910732085268e+00
+4.139052155837465996e+00
+4.390305777018236988e+00
+4.642172339120373081e+00
+4.894678666149941293e+00
+5.147269052361176378e+00
+5.398003206656323272e+00
+5.647798235368590269e+00
+5.900534302989873581e+00
+6.151460297009074374e+00
+6.402630176204441526e+00
+6.655964674748322096e+00
+6.908782583105771202e+00
+7.159060853508338873e+00
+7.409414762591408632e+00
+7.659227149364629561e+00
+7.911559474606142395e+00
+8.163750941790897997e+00
+8.415757210439384295e+00
+8.667638839289207198e+00
+8.918801124864609520e+00
+9.170480108345369530e+00
+9.420566033467249412e+00
+9.671642335271435797e+00
+9.923426157758477473e+00
+1.017396902749181820e+01
+1.042533700753922155e+01
+1.067702948358784276e+01
+1.092969337256462303e+01
+1.118211872033766419e+01
+1.143346516925797296e+01
+1.168350508557724687e+01
+1.193379937775929100e+01
+1.218553055762879112e+01
+1.243711650801278878e+01
+1.268867268647359303e+01
+1.294075027177769144e+01
+1.319267287020158363e+01
+1.344470571782862756e+01
+1.369480150283709641e+01
+1.394591277066507828e+01
+1.419764645764527167e+01
+1.444766774168416212e+01
+1.469906835825429070e+01
+1.495059118111036511e+01
+1.520261612590052991e+01
+1.545415944503929495e+01
+1.570456072545780124e+01
+1.595705800128773966e+01
+1.620829297830731619e+01
+1.645966591028190962e+01
+1.671125708846596325e+01
+1.696160259490052624e+01
+1.721370276108077135e+01
+1.746543424102250341e+01
+1.771598473288305442e+01
+1.796653767411944003e+01
+1.821765658501893270e+01
+1.847053788120423334e+01
+1.872257696809484173e+01
+1.897394097105746269e+01
+1.922542454726159988e+01
+1.947607703355479103e+01
+1.972678839839982601e+01
+1.997802698857814363e+01

README.md

@@ -0,0 +1,71 @@
+# CosmoGraphNet
+
+[![DOI](https://zenodo.org/badge/DOI/10.5281/zenodo.xxxx.svg)](https://doi.org/10.5281/zenodo.xxxx) [![arXiv](https://img.shields.io/badge/arXiv-2204.xxxxx-B31B1B.svg)](http://arxiv.org/abs/2204.xxxxx)
+
+Graph Neural Networks to predict the cosmological parameters or the galaxy power spectrum from galaxy catalogs.
+
+A graph is created from a galaxy catalog with information about the 3D position and intrinsic galactic properties. A Graph Neural Network is then applied to infer the cosmological parameters or the galaxy power spectrum. Galaxy catalogs extracted from the [CAMELS](https://camels.readthedocs.io/en/latest/index.html) hydrodynamic simulations, specially suited for Machine Learning purposes. Neural nets architectures are defined making use of the package [PyTorch-geometric](https://pytorch-geometric.readthedocs.io/en/latest/).
+
+See the paper [arXiv:2204.xxxxx](https://arxiv.org/abs/2204.xxxxx) for more details.
+
+<img src="visualize_graph_10.png" width="500">
+
+
+## Scripts
+
+Here is a brief description of the codes included:
+
+- `main.py`: main driver to train and test the network.
+
+- `hyperparameters.py`: script with the definition of the hyperparameters employed by the networks.
+
+- `crosstest.py`: tests a pre-trained model.
+
+- `hyperparams_optimization.py`: optimize the hyperparameters using `optuna`.
+
+- `ps_test.py`: tests the power spectrum neural networks in point distributions with different clustering properties.
+
+- `visualize_graphs.py`: display graphs from galaxy catalogs in 2D or 3D.
+
+
+The folder `Source` contains scripts with auxiliary routines:
+
+* `constants.py`: basic constants and initialization.
+
+* `load_data.py`: contains routines to load data from simulation files.
+
+* `plotting.py`: includes functions for displaying the results from the neural nets.
+
+* `metalayer.py`: includes the definition of the Graph Neural Networks architecture.
+
+* `training.py`: includes routines for training and testing the net.
+
+
+## Requisites
+
+The libraries required for training the models and compute some statistics are:
+* `numpy`
+* `pytorch`
+* `pytorch-geometric`
+* `matplotlib`
+* `scipy`
+* `sklearn`
+* `optuna` (only for optimization in `hyperparams_optimization.py`)
+
+
+## Usage
+
+These are some advices to employ the scripts described above:
+1. To perform a search of the optimal hyperparameters, run `hyperparams_optimization.py`.
+2. To train a model with a given set of parameters defined in `hyperparameters.py`, run `main.py`. The hyperparameters currently present in `hyperparameters.py` correspond to the best optimal values for each suite when all galactic features are employed (see the paper).
+3. Once a model is trained, run `crosstest.py` to test in the training simulation suite and cross test it in the other one included in CAMELS (IllustrisTNG and SIMBA).
+
+
+## Citation
+
+If you use the code, please link this repository, and cite [arXiv:2204.xxxxx](https://arxiv.org/abs/2204.xxxxx) and the DOI [10.5281/zenodo.xxxx](https://doi.org/10.5281/zenodo.xxxx).
+
+
+## Contact
+
+Feel free to contact me at <[email protected]> for comments, questions and suggestions.

Source/constants.py

@@ -0,0 +1,53 @@
+#----------------------------------------------------------------------
+# List of constants and some common functions
+# Author: Pablo Villanueva Domingo
+# Last update: 10/11/21
+#----------------------------------------------------------------------
+
+import numpy as np
+import torch
+import os
+import random
+
+# Random seeds
+torch.manual_seed(12345)
+np.random.seed(12345)
+random.seed(12345)
+
+# use GPUs if available
+if torch.cuda.is_available():
+    print("CUDA Available")
+    device = torch.device('cuda')
+else:
+    print('CUDA Not Available')
+    device = torch.device('cpu')
+
+#--- PARAMETERS AND CONSTANTS ---#
+
+# Reduced Hubble constant
+hred = 0.7
+
+# Root path for simulations
+simpathroot = "/projects/QUIJOTE/CAMELS/Sims/"
+
+# Box size in comoving kpc/h
+boxsize = 25.e3
+
+# Validation and test size
+valid_size, test_size = 0.15, 0.15
+
+# Batch size
+batch_size = 25
+
+
+
+# Number of k bins in the power spectrum
+ps_size = 79
+
+#--- FUNCTIONS ---#
+
+
+# Choose color depending on the CAMELS simulation suite
+def colorsuite(suite):
+    if suite=="IllustrisTNG":   return "purple"
+    elif suite=="SIMBA":            return "deepskyblue"