add argument for CUDA version and reformat

Dockerfile

@@ -4,6 +4,9 @@ ARG MPI=openmpi
 # Select Platform: core, full, pycuda or cupy
 ARG PLATFORM=cupy
 
+# Select CUDA version
+ARG CUDAVERSION=12.4
+
 # Pull from mambaforge and install XML and ssh
 FROM condaforge/mambaforge as base
 ENV DEBIAN_FRONTEND=noninteractive
@@ -26,26 +29,46 @@ RUN mamba env update -n base -f dependencies.yml
 
 # Pull from MPI build and install accelerate/pycuda dependencies
 FROM mpi as pycuda
+ARG CUDAVERSION
 COPY ./ptypy/accelerate/cuda_pycuda/dependencies.yml ./dependencies.yml
 COPY ./cufft/dependencies.yml ./dependencies_cufft.yml
-RUN mamba env update -n base -f dependencies.yml && mamba env update -n base -f dependencies_cufft.yml
+RUN mamba install cuda-version=${CUDAVERSION} && \
+    mamba env update -n base -f dependencies.yml && \
+    mamba env update -n base -f dependencies_cufft.yml
 
 # Pull from MPI build and install accelerate/cupy dependencies
 FROM mpi as cupy
+ARG CUDAVERSION
 COPY ./ptypy/accelerate/cuda_cupy/dependencies.yml ./dependencies.yml
 COPY ./cufft/dependencies.yml ./dependencies_cufft.yml
-RUN mamba env update -n base -f dependencies.yml && mamba env update -n base -f dependencies_cufft.yml
+RUN mamba install cuda-version=${CUDAVERSION} && \
+    mamba env update -n base -f dependencies.yml && \
+    mamba env update -n base -f dependencies_cufft.yml
 
-# Pull from existing image with ptypy dependencies and set up testing
-FROM ${PLATFORM} as runtime
+# Pull from platform specific image and install ptypy 
+FROM ${PLATFORM} as build
 COPY pyproject.toml setup.py ./
 COPY ./scripts ./scripts
 COPY ./templates ./templates
 COPY ./benchmark ./benchmark
 COPY ./cufft ./cufft
 COPY ./ptypy ./ptypy
 RUN pip install .
-RUN if [ "$PLATFORM" = "pycuda" ] || [ "$PLATFORM" = "cupy" ] ; then pip install ./cufft ; fi
+
+# For core/full build, no post processing needed
+FROM build as core-post
+FROM build as full-post
+
+# For pycuda build, install filtered cufft
+FROM build as pycuda-post
+RUN pip install ./cufft
+
+# For pycuda build, install filtered cufft
+FROM build as cupy-post
+RUN pip install ./cufft
+
+# Platform specific runtime container
+FROM ${PLATFORM}-post as runtime
 
 # Run PtyPy run script as entrypoint
 ENTRYPOINT ["ptypy.run"]