Add code

Author: marvinsxtr

.github/workflows/docker-publish.yaml

@@ -32,6 +32,9 @@ jobs:
       - name: Checkout repository
         uses: actions/checkout@v4
 
+      - name: Clean up disk space
+        uses: jlumbroso/free-disk-space@main
+
       # Install the cosign tool except on PR
       # https://github.com/sigstore/cosign-installer
       - name: Install cosign
@@ -77,30 +80,29 @@ jobs:
           cache-from: type=gha
           cache-to: type=gha,mode=max
 
-      # Sign the resulting Docker image digest except on PRs.
-      # This will only write to the public Rekor transparency log when the Docker
-      # repository is public to avoid leaking data.  If you would like to publish
-      # transparency data even for private images, pass --force to cosign below.
-      # https://github.com/sigstore/cosign
-      - name: Sign the published Docker image
-        if: ${{ github.event_name != 'pull_request' }}
-        env:
-          # https://docs.github.com/en/actions/security-guides/security-hardening-for-github-actions#using-an-intermediate-environment-variable
-          TAGS: ${{ steps.meta.outputs.tags }}
-          DIGEST: ${{ steps.build-and-push.outputs.digest }}
-        # This step uses the identity token to provision an ephemeral certificate
-        # against the sigstore community Fulcio instance.
-        run: echo "${TAGS}" | xargs -I {} cosign sign --yes {}@${DIGEST}
+      - name: Free disk space after Docker build
+        run: |
+          docker buildx ls --format '{{.Name}}' | grep -v default | xargs -I {} docker buildx rm {} 2>/dev/null || true
+          docker system prune -af --volumes
+          docker builder prune -af
+
+          rm -rf ~/.docker/buildx || true
+          rm -rf /tmp/docker-actions-toolkit-*/* 2>/dev/null || true
+
+          df -h /
+
+      - id: lower-repo
+        name: Repository to lowercase
+        run: |
+          echo "repository=${GITHUB_REPOSITORY@L}" >> $GITHUB_OUTPUT
 
       # Build and push .sif files for Apptainer
-      - name: Clean up disk space
-        uses: jlumbroso/free-disk-space@main
       - name: Setup Apptainer
         uses: eWaterCycle/setup-apptainer@v2
       - name: Build and push Apptainer
         env:
           TAGS: ${{ steps.meta.outputs.tags }}
         run: |
            echo ${{ secrets.GITHUB_TOKEN }} | apptainer registry login -u ${{ secrets.GHCR_USERNAME }} --password-stdin docker://ghcr.io
-           apptainer build container.sif docker://${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:latest
+           apptainer build container.sif docker://${{ env.REGISTRY }}/${{ steps.lower-repo.outputs.repository }}:latest
            echo "${TAGS}" | xargs -I {} apptainer push container.sif oras://{}-sif

.github/workflows/pytest.yaml

@@ -16,7 +16,7 @@ jobs:
           cache-dependency-glob: "uv.lock"
 
       - name: Install the project
-        run: uv sync --all-extras --dev
+        run: uv sync --all-extras --all-groups
 
       - name: Run tests
         run: uv run pytest

.github/workflows/ruff.yaml

@@ -7,4 +7,4 @@ jobs:
       - uses: actions/checkout@v4
       - uses: astral-sh/ruff-action@v3
         with:
-          version: 0.8.0
+          version: 0.8.0

.gitignore

@@ -161,5 +161,14 @@ cython_debug/
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
 
-outputs
-wandb
+*.csv
+
+wandb
+lightning_logs
+perturbench_data
+
+/outputs
+/meta_flow_matching
+/condot
+/logs
+/data

Dockerfile

@@ -1,10 +1,8 @@
 FROM --platform=linux/amd64 python:3.12-slim AS linux-base
 
 # Utilities
-RUN apt-get update && apt-get upgrade -y
-RUN apt-get install -y --no-install-recommends build-essential \
-    sudo curl git htop less rsync screen vim nano wget ca-certificates \
-    openssh-client zsh
+RUN apt-get update && apt-get upgrade -y && apt-get install -y --no-install-recommends build-essential \
+    sudo curl git htop less rsync screen vim nano wget ca-certificates openssh-client zsh procps psmisc
 
 # Download and install VS Code Server CLI
 RUN wget -O /tmp/vscode-server-cli.tar.gz "https://update.code.visualstudio.com/latest/cli-linux-x64/stable" && \
@@ -31,6 +29,8 @@ ENV UV_LINK_MODE=copy
 ENV UV_PYTHON=python3.12
 ENV PATH="$UV_PROJECT_ENVIRONMENT/bin:$PATH"
 ENV PYTHONPATH="/srv/repo:$PYTHONPATH"
+# See https://github.com/jax-ml/jax/issues/29260
+ENV LD_LIBRARY_PATH=/venv/lib/python3.12/site-packages/nvidia/cuda_nvrtc/lib:$LD_LIBRARY_PATH
 
 # Install uv
 COPY --from=ghcr.io/astral-sh/uv:0.6.6 /uv /usr/local/bin/uv
@@ -39,4 +39,4 @@ COPY --from=ghcr.io/astral-sh/uv:0.6.6 /uv /usr/local/bin/uv
 RUN --mount=type=cache,target=/root/.cache/uv \
     --mount=type=bind,source=uv.lock,target=uv.lock \
     --mount=type=bind,source=pyproject.toml,target=pyproject.toml \
-    uv sync --frozen --no-install-project --no-dev
+    uv sync --frozen --no-install-project --all-groups

LICENSE

@@ -1,6 +1,6 @@
 MIT License
 
-Copyright (c) 2025 Marvin Sextro
+Copyright (c) 2026 Marvin Sextro
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal