README.md

---
license: cc-by-sa-4.0
datasets:
- Smith42/galaxies
- Smith42/galaxies_metadata
- Smith42/galaxies_embeddings
tags:
- astronomy
- images
- huggingscience
- science
---
<center>
<img src="assets/shoggoth_telescope_sticker_2.png" alt="astroPT_shoggoth" width="300px"/>
</center>

# astroPTv2.0: a Large Observation Model for Astronomy

Here we have the model files for the astroPT project, the code to run inference
with these models is found here:
[https://github.com/smith42/astropt](https://github.com/smith42/astropt)

You will find the fully trained models (pretrained on 8.6 million galaxies) in
folders labelled with the model parameter count in the `astropt` directory.

Unlike the older models which were trained on the "image" column in 
[smith42/galaxies](https://huggingface.co/datasets/Smith42/galaxies), 
these models are trained on the "cropped" galaxies from the "image_crop"
column. Those galaxies have been cropped and zoomed so that they take up the 
majority of each image before uploading.

We get some promising scaling on this new dataset, see below:

<center>
<img src="assets/scaling_law.png" alt="scaling_law" width="300px"/>
</center>

## Usage

To use these models in anger you can `pip install astropt` and run the following code:

```python
from astropt.model_utils import load_astropt
from astropt.local_datasets import GalaxyImageDataset

from datasets import load_dataset # for Smith42/galaxies

import torch
import numpy as np
from functools import partial
from torch.utils.data import DataLoader
from torchvision import transforms

# boilerplate to preprocess galaxy images
def normalise(x):
    std, mean = torch.std_mean(x, dim=1, keepdim=True)
    return (x - mean) / (std + 1e-8)

def data_transforms():
    return transforms.Compose([transforms.Lambda(normalise)])

def _process_galaxy_wrapper(idx, func):
    """This function ensures that the image is tokenised in the same way as the pre-trained model is expecting"""
    galaxy = func(
        torch.from_numpy(np.array(idx["image"]).swapaxes(0, 2)).to(float)
    ).to(torch.float)
    galaxy_positions = torch.arange(0, len(galaxy), dtype=torch.long)
    return {
        "images": galaxy,
        "images_positions": galaxy_positions,
    }

# for 095M parameter model, 015M and 850M models are also available:
model = load_astropt("Smith42/astroPT_v2.0", path="astropt/095M")

galproc = GalaxyImageDataset(
  None,
  spiral=True,
  transform={"images": data_transforms()},
  modality_registry=model.modality_registry
)

ds = (
  load_dataset("Smith42/galaxies", split="test", revision="v2.0", streaming=True)
  .select_columns("image")
  .map(partial(_process_galaxy_wrapper, func=galproc.process_galaxy))
  .with_format("torch")
)

dl = iter(DataLoader(ds, batch_size=128, num_workers=32))

zs = []
for B in dl:
    zs.append(model.generate_embeddings(B)["images"].detach().numpy())
zs = np.concatenate(zs)

# do cool stuff with zs...
```


## Updates and community

AstroPT is an open-to-all UniverseTBD project. Please join the [UniverseTBD](https://universetbd.org) Discord for updates: [https://discord.gg/MNEVegvfJq](https://discord.gg/MNEVegvfJq)