Merge pull request #46 from UM-Bridge/gs2

Adds GS2 to benchmark models

.github/workflows/model-gs2.yml

@@ -0,0 +1,49 @@
+name: model-gs2
+
+on:
+  push:
+    branches:
+      - 'main'
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    steps:
+      -
+        name: Checkout
+        uses: actions/checkout@v2
+      -
+        name: Set up QEMU
+        uses: docker/setup-qemu-action@v1
+      -
+        name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v2
+      -
+        name: Login to DockerHub
+        uses: docker/login-action@v1
+        with:
+          username: ${{ secrets.DOCKERHUB_USERNAME }}
+          password: ${{ secrets.DOCKERHUB_TOKEN }}
+      -
+        name: Build and push
+        uses: docker/build-push-action@v3
+        with:
+          context: models/gs2/
+          push: true
+          tags: linusseelinger/model-gs2:latest
+
+  test:
+    runs-on: ubuntu-latest
+    needs: build
+
+    services:
+      model:
+        image: linusseelinger/model-gs2:latest
+        ports:
+        - 4242:4242
+
+    steps:
+       -
+        name: Validate
+        run: |
+          docker run --network=host -e model_host=http://localhost:4242 linusseelinger/testing-protocol-conformity-current:latest

models/gs2/Dockerfile

@@ -0,0 +1,52 @@
+FROM mpioperator/openmpi
+LABEL maintainer="[email protected]"
+LABEL version="1.0"
+LABEL description="This is Docker Image for running GS2 with GNU compilers on \
+Ubuntu 20.04. This uses MPICH instead of OpenMPI to avoid issues with running \
+as root."
+
+# Disable Prompt During Packages Installation
+ARG DEBIAN_FRONTEND=noninteractive
+
+RUN apt update \ 
+	&& apt install -y \
+	gfortran \
+	g++ \
+	make \
+	wget \
+	git \
+	vim \
+	cmake \
+	rsync \
+	curl \
+	python3 \
+	python3.8-venv \
+	python3-pip \
+	mpich \
+	libnetcdf-dev \
+	libnetcdff-dev \
+	libfftw3-dev \
+        python-is-python3
+
+RUN echo 'alias python=python3' >> $HOME/.bashrc
+RUN echo 'alias pip=pip3' >> $HOME/.bashrc
+
+RUN echo 'export GK_SYSTEM=gnu_ubuntu' >> $HOME/.bashrc
+
+RUN pip install ford netCDF4==1.6.2
+
+RUN git clone https://[email protected]/gyrokinetics/gs2.git /usr/gs2
+
+WORKDIR /usr/gs2
+
+RUN cd /usr/gs2
+
+ENV GK_SYSTEM=gnu_ubuntu
+
+RUN echo ${GK_SYSTEM}
+
+RUN make -IMakefiles depend
+
+RUN make -IMakefiles gs2
+
+#RUN make -IMakefiles tests

models/gs2/README.md

@@ -0,0 +1,56 @@
+# GS2 Fusion Plasma Simulation
+
+## Overview
+This model uses [GS2](https://gyrokinetics.gitlab.io/gs2/index.html) to study plasma in a spherical tokamak; the simulation has been configured to use the gyrokinetic framework. It aims to study various modes in the plasma due to microinstabilities. As opposed to finding all available unstable modes, the current setup terminates once an unstable mode is found. Otherwise, it will continue running until reaching a fixed timestep. Therefore, the runtime varies depending on the input parameters.
+
+## Authors
+- [Chung Ming Loi](mailto:[email protected])
+
+
+## Run
+```
+docker run -it -p 4242:4242 linusseelinger/model-gs2:latest
+```
+
+## Properties
+
+Model | Description
+---|---
+forward | Plasma simulation with the gyrokinetic model
+
+### forward
+Mapping | Dimensions | Description
+---|---|---
+input | [2] | [`tprim`: normalised inverse temperature gradient, `vnewk`: normalised species-species collisionality frequency]. Both are set for electrons only. 
+output | [3] | [Electron heat flux, electric field growth rate, electric field mode frequency]
+
+Feature | Supported
+---|---
+Evaluate | True
+Gradient | False
+ApplyJacobian | False
+ApplyHessian | False
+
+Config | Type | Default | Description
+---|---|---|---
+None | | |
+
+## Mount directories
+Mount directory | Purpose
+---|---
+None |
+
+## Source code
+
+[Model sources here.](https://github.com/UM-Bridge/benchmarks/tree/gs2/models/gs2)
+
+## Description
+GS2 is developed to study low frequency turbulence in magnetised plasma. In this benchmark, it solves the gyrokinetic Vlasov-Maxwell system of equations that are widely adopted to describe the turbulent component of the electromagnetic fields and particle distribution functions of the species present in a plasma. The problem is five-dimensional and takes the form 
+
+$$
+\frac{\partial F_a}{\partial t} + \frac{\partial \vec{X}}{\partial t} \cdot \nabla F_a + \frac{\partial v_{\parallel}}{\partial t} \frac{\partial F_a}{\partial v_{\parallel}} = 0, 
+$$
+
+where $F_a = F_a(\vec{x}, v_{\parallel}, \mu)$ is the 5D phase space gyrocenter distribution function for species $a$, $\frac{\partial \vec{X}}{\partial t} = \vec{v}$ is the particle velocity, $\vec{x}$ is the three dimensional vector describing the guiding centre of a particle, $v_{\parallel}$ is the velocity along the magnetic field line and $\mu = v^2_{\perp}/2B$ is the magnetic moment, where $v_{\perp}$ is the velocity perpendicular to the magnetic field and $B$ is the magnetic flux density.
+
+The above text was taken from the paper by [Hornsby et al. (2023)](https://doi.org/10.48550/arXiv.2309.09785) which contains a more detailed description of the problem.

models/gs2/fast.in

@@ -0,0 +1,170 @@
+&kt_grids_knobs
+    grid_option = 'single'
+/
+
+&kt_grids_single_parameters
+    aky = 0.7852773972394589
+    theta0 = 0.0
+/
+
+&theta_grid_parameters
+    ntheta = 128
+    nperiod = 9
+    akappa = 2.66
+    akappri = -0.25
+    rhoc = 0.67
+    tri = 0.3400000000000001
+    tripri = 0.25
+    shat = 0.87
+    qinp = 8.903370045724728
+    r_geo = 1.84
+    rmaj = 1.84
+    shift = -0.44
+    geotype = 0
+/
+
+&theta_grid_knobs
+    equilibrium_option = 'eik'
+/
+
+&theta_grid_eik_knobs
+    irho = 2
+    iflux = 0
+    s_hat_input = 3.1173933836235315
+    beta_prime_input = -0.64
+    local_eq = .true.
+    bishop = 4
+    equal_arc = .false.
+    writelots = .true.
+/
+
+&le_grids_knobs
+    ngauss = 8
+    negrid = 16
+    bouncefuzz = 1e-08
+/
+
+&dist_fn_knobs
+    adiabatic_option = 'iphi00=2'
+    opt_source = .true.
+    def_parity = .true.
+    even = .false.
+/
+
+&fields_knobs
+    field_option = 'implicit'
+    response_dir = 'response'
+/
+
+&knobs
+    fphi = 1.0
+    fapar = 1.0
+    fbpar = 0.0
+    delt = 0.02
+    nstep = 80000
+    wstar_units = .true.
+/
+
+&layouts_knobs
+    layout = 'xyles'
+/
+
+&collisions_knobs
+    collision_model = 'default'
+/
+
+&species_knobs
+    nspec = 3
+/
+
+&species_parameters_1
+    z = 1.0
+    mass = 1.0
+    dens = 0.5
+    temp = 1.0593220338983051
+    tprim = 0.0
+    fprim = 0.1450502618121151
+    uprim = 0.0
+    vnewk = 0.00039999999999999996
+    type = 'ion'
+    bess_fac = 1.0
+/
+
+&dist_fn_species_knobs_1
+    fexpr = 0.48
+    bakdif = 0.05
+/
+
+&species_parameters_2
+    z = 1.0
+    mass = 1.4975
+    dens = 0.5
+    temp = 1.0593220338983051
+    tprim = 0.0
+    fprim = 0.1450502618121151
+    uprim = 0.0
+    vnewk = 0.0003268711385781747
+    type = 'ion'
+    bess_fac = 1.0
+/
+
+&dist_fn_species_knobs_2
+    fexpr = 0.48
+    bakdif = 0.05
+/
+
+&species_parameters_3
+    z = -1
+    mass = 0.0002724
+    dens = 1.0
+    temp = 1.0
+    tprim = 4.98536245337014
+    fprim = 2.0329154559480602
+    uprim = 0.0
+    vnewk = 0.011195269013375201
+    type = 'electron'
+    bess_fac = 1.0
+/
+
+&dist_fn_species_knobs_3
+    fexpr = 0.48
+    bakdif = 0.05
+/
+
+&init_g_knobs
+    ginit_option = 'default_odd'
+    chop_side = .false.
+    phiinit = 1e-05
+    restart_dir = 'restart'
+/
+
+&gs2_diagnostics_knobs
+    write_ascii = .true.
+    write_omega = .true.
+    write_final_fields = .true.
+    write_fields = .true.
+    write_phi_over_time = .true.
+    write_bpar_over_time = .true.
+    write_apar_over_time = .true.
+    write_nl_flux_dist = .true.
+    write_fluxes = .true.
+    write_nl_flux = .true.
+    nwrite = 30
+    navg = 10
+    omegatol = 0.0001
+    omegatinst = 500.0
+    nsave = 5000
+    save_for_restart = .true.
+    write_final_epar = .true.
+/
+
+&parameters
+    beta = 0.29798123734456977
+    tite = 1.0
+    zeff = 1.0
+/
+
+&diagnostics_config
+    nwrite = 100000000
+/
+

models/gs2/server.py

@@ -0,0 +1,57 @@
+import umbridge
+import json
+import os
+import csv
+from pyrokinetics import Pyro
+import numpy as np
+from datetime import datetime
+from fileinput import FileInput
+
+
+
+class GS2Model(umbridge.Model):
+    def __init__(self):
+        super().__init__("forward")
+
+    def get_input_sizes(self, config):
+        return [2] 
+
+    def get_output_sizes(self, config):
+        return [3] 
+
+    def __call__(self, parameters, config):
+        input_file = "fast.in" # Select input file
+        os.system("mkdir restart") # GS2 needs this folder otherwise will fail
+        with FileInput(files=input_file, inplace=True) as file:
+            checkpoint = 0
+            for line in file:
+                if "&species_parameters_3" in line:
+                    checkpoint += 1
+                if "tprim" in line and checkpoint == 1:
+                    line = f"    tprim = {parameters[0][0]}"
+                if "vnewk" in line and checkpoint == 1:
+                    line = f"    vnewk = {parameters[0][1]}"
+                print(line.rstrip()) # FileInput redirects stdout to file
+
+        # Run the model 
+        mpirank = config.get("ranks", 32)
+        os.system(f"mpirun -np {mpirank} /usr/gs2/bin/gs2 {input_file}")
+
+        # Read results using Pyrokinetics package and print output
+        gs2_input = input_file
+        pyro = Pyro(gk_file=gs2_input, gk_code="GS2")
+        pyro.load_gk_output()
+        data = pyro.gk_output
+        print(data)
+        heat = data["heat"].sel(species="electron", field='phi').isel(time=-1)
+        growth_rate = data["growth_rate"].isel(time=-1)
+        mode_freq = data["mode_frequency"].isel(time=-1)
+        output = [[heat.to_numpy().squeeze().tolist(), growth_rate.to_numpy().squeeze().tolist(), mode_freq.to_numpy().squeeze().tolist()]]
+        return output 
+
+    def supports_evaluate(self):
+        return True
+
+model = GS2Model()
+
+umbridge.serve_models([model], 4242)