Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

StarNet -- An Inference method for Population III star formation and feedback #377

Draft
wants to merge 41 commits into
base: main
Choose a base branch
from

Conversation

WillHicks96
Copy link
Contributor

@WillHicks96 WillHicks96 commented Mar 20, 2024

This WIP PR builds off of PR #274, and adds the ability to call pre-trained StarNet models using the C++ frontend of PyTorch (https://pytorch.org/cppdocs/frontend.html) to predict the locations of Pop III stars within an evolving cosmological simulation, as well as the cumulative feedback effects once all stars within a given Pop III association die.

There are a number of changes:

  1. Users can now optionally link Enzo-E to the Torch library at compile time by specifying -DUSE_TORCH=ON and -DTorch_ROOT=${TORCH_DIR} CMAKE parameters.
  2. I have added on to EnzoMethodInference under the comment that reads "ADD DEEP LEARNING INFERENCE HERE". The added code here deserializes pre-trained StarNet checkpoints, accesses the relevant fields for making predictions, rescales the fields, packages them into a torch::Tensor object, forwards the torch::Tensor object to the pre-trained models, interprets the results to identify cells within an inference block that are flagged for star formation, computes population statistics, and finally predicts the size of the conglomerate supernova remnant and metal yields. A tracer particle of type popIIIremnant is added onto the mesh at the center of mass of each predicted star-forming region, which contains information about the population statistics, yields, etc. The thought here is that these particles can be analyzed in post to connect Pop III metals with their progenitor stellar populations.
  3. So as not to modify EnzoMethodInference too much to the point that it can't be used for other inference models in the future, I have added the EnzoMethodFBNetDeposit method to handle the actual deposition Pop III metals and updates to various fields associated with the feedback. Metals are deposited into separate fields based on supernova type (Type II SNe, HNe, or PISNe).
  4. I have added the EnzoStarNetUtils and EnzoObjectFeedbackSphere classes to help out with the calculations and data manipulation mentioned above.

StarNet was originally developed and trained in classic Enzo using its inline python capabilities (for more info, see Wells & Norman (2021), Wells & Norman (2022a), and Wells & Norman (2022b)).

PR Checklist

This is a WIP PR is contingent on:

  • Streamlining the process for installing the C++ frontend of Torch. The way I got it to work originally was kind of hacky, and I'm honestly scared to touch it because I don't want my installation to break before I graduate.
  • PR Inference array #274 itself is a WIP PR, so that one needs to merge before this one.
  • documentation
  • tests -- since PR also links a new external library, I'm not sure exactly how the automated compilation/testing suite on GitHub should be updated for the purposed of running unit tests. It might make sense to split off the CMAKE changes for linking to the Torch library to its own standalone PR.
  • General code cleaning

Here's is an image from a 512^3 simulation I ran that shows StarNet in action (publication in progress!):

projections_with_zooms

This is at z=11. Pop III feedback is the only source of metal enrichment in this simulation, so the metal bubbles are all coming from StarNet. The box length is 5 Mpc. The zoomed-in images show density and metallicity slices of the two most massive halos (Mvir = 1e9 Msun roughly).

William Hicks and others added 30 commits January 4, 2023 16:33
Not working fully -- get NaNs and segfaults
…e in StarFind and FBNet, along with txt files containing distribution functions relevent to FBNet that were generated using the PHX simulations
…lso introduced EnzoObjectFeedbackSphere class to store positions, radii, metal_yields, etc. for SN remnants
…e feedbackspheres are only sent to blocks that contribute to the level array
Multiplying radius by factor of 0.2 to correct for overestimations from PHX
Tweaked metallicity scaling
William Hicks and others added 10 commits July 12, 2023 14:35
Not seeding RNG before every call to rand() anymore so that it doesn't just return the same number over and over again
Added another filtering criterion that checks the mean overdensity of the block, rather than cell-wise overdensity
…der, regardless of the order of Group:inference:field_list
…es to EnzoObjectFeedbackSphere and remnant particles. Fixed stellar mass sum
@mabruzzo
Copy link
Contributor

I'm definitely happy to help with CMake.

  • Way back when you mentioned this all to me, I had looked into streamlining things.
  • Since then, I've learned more now about CMake. It's been a while since I looked at libtorch, but I now think that streamlining this may actually be fairly straightforward.

The key point is that libtorch installations come with a cmake "config-file". We should be able to use that (rather than creating a FindTorch.cmake file). If we do that properly, the installation should just work™. (As I say this, I think I had started down this path before and the "doing it properly" part was the tricky bit -- but now that I know a lot more about these config-files I'm confident I can do it)

@jwise77 jwise77 self-requested a review March 21, 2024 15:56
@jwise77 jwise77 requested a review from bwoshea April 19, 2024 14:02
@bwoshea
Copy link
Contributor

bwoshea commented Jan 16, 2025

@WillHicks96 would you mind bringing this branch up-to-date with the main branch of Enzo-E? There are a bunch of conflicts that have to be resolved. Thank you!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants