Proper setup of PyTorch versions in Conda based environments (#1492)

* Limit permissible PyTorch versions in Conda

Conda would typically install the latest version of any dependency,
but we don't support the latest version of PyTorch yet. Keep version
aligned with settings in setup.py for now.

* Pull PyTorch from pytorch Conda channel

The desired PyTorch version is not available from default channels, such
that we need to set the pytorch channel explicitly.

* CUDA environments for Conda and Python update

Use Python 3.11, which is the latest compatible with Heat and Pytorch.

Introduce CUDA 11.8 and 12.1 environments for Conda.

* Update quick_start.md

* Update quick_start.md

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update quick_start.md

---------

Co-authored-by: Fabian Hoppe <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Claudia Comito <[email protected]>
Co-authored-by: Michael Tarnawa <[email protected]>

quick_start.md

@@ -11,7 +11,7 @@ The Heat conda build includes all dependencies including OpenMPI.
 ```shell
 conda create --name heat_env
 conda activate heat_env
-conda -c conda-forge heat
+conda install -c conda-forge heat
 ```
 
 [Test](#test) your installation.
@@ -78,21 +78,34 @@ Local torch tensor on rank  1 :  tensor([5, 6, 7, 8, 9], dtype=torch.int32)
 
 3. [Fork](https://docs.github.com/en/get-started/quickstart/contributing-to-projects) or, if you have write access, clone the [Heat repository](https://github.com/helmholtz-analytics/heat).
 
-4. Create a virtual environment `heat_dev` with all dependencies via [heat_dev.yml](https://github.com/helmholtz-analytics/heat/blob/main/scripts/heat_dev.yml). Note that `heat_dev.yml` does not install Heat.
+4. **Setting up a dev-environment with CONDA:** Create a virtual environment `heat_dev` with all dependencies via [scripts/heat_dev.yml](https://github.com/helmholtz-analytics/heat/blob/main/scripts/heat_dev.yml). Note that `scripts/heat_dev.yml` does not install Heat.
 
     ```
-    conda env create -f heat_dev.yml
+    conda env create -f scripts/heat_dev.yml
     conda activate heat_dev
     ```
+    Note that in case you want to use a GPU while developing on your local machine, you need to set up a CUDA environment by using `scripts/heat_dev_cuda11.yml`for CUDA 11 or `scripts/heat_dev_cuda12.yml`for CUDA 12, respectively, instead of `scripts/heat_dev.yml`.
 
-5. In the `/heat` directory of your local repo, install the [pre-commit hooks]( https://pre-commit.com/):
+    **Setting up a dev-environment with PIP:** Create a virtual environment `heatenv` with `python -m venv <path_to_store_venvs>/heatenv`, and activate it by `source <path_to_store_venvs>/heatenv/bin/activate`. Then clone the Heat-repo from GitHub by
+
+    ```
+    git clone https://github.com/helmholtz-analytics/heat.git
+    ```
+
+     go to the Heat-folder (`cd heat`), and install (in editable fashion "`-e`") by
+
+    ```
+    pip install -e '.[hdf5, netcdf]'
+    ```
+
+6. In the `/heat` directory of your local repo, install the [pre-commit hooks]( https://pre-commit.com/):
 
     ```
     cd $MY_REPO_DIR/heat/
     pre-commit install
     ```
 
-6. Write and run (locally) [unit tests](https://docs.python.org/3/library/unittest.html) for any change you introduce. Here's a sample of our [test modules](https://github.com/helmholtz-analytics/heat/tree/main/heat/core/tests).
+7. Write and run (locally) [unit tests](https://docs.python.org/3/library/unittest.html) for any change you introduce. Here's a sample of our [test modules](https://github.com/helmholtz-analytics/heat/tree/main/heat/core/tests).
 
     Running all unit tests locally, e.g. on 3 processes:
 
@@ -125,6 +138,6 @@ Local torch tensor on rank  1 :  tensor([5, 6, 7, 8, 9], dtype=torch.int32)
     mpirun --tag-output -n 3 python -m unittest -vf
     ```
 
-7. After [making and pushing](https://docs.github.com/en/get-started/quickstart/contributing-to-projects#making-and-pushing-changes) your changes, go ahead and [create a Pull Request](https://docs.github.com/en/get-started/quickstart/contributing-to-projects#making-a-pull-request). Make sure you go through the Due Diligence checklist (part of our PR template). Consider [allowing us to edit your branch](https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/allowing-changes-to-a-pull-request-branch-created-from-a-fork#enabling-repository-maintainer-permissions-on-existing-pull-requests) for a smoother review process.
+8. After [making and pushing](https://docs.github.com/en/get-started/quickstart/contributing-to-projects#making-and-pushing-changes) your changes, go ahead and [create a Pull Request](https://docs.github.com/en/get-started/quickstart/contributing-to-projects#making-a-pull-request). Make sure you go through the Due Diligence checklist (part of our PR template). Consider [allowing us to edit your branch](https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/allowing-changes-to-a-pull-request-branch-created-from-a-fork#enabling-repository-maintainer-permissions-on-existing-pull-requests) for a smoother review process.
 
     ## Thank you so much for your time!

scripts/heat_dev.yml

@@ -3,12 +3,12 @@ channels:
   - conda-forge
   - defaults
 dependencies:
-  - python=3.10
+  - python=3.11
   - openmpi
   - mpi4py
   - h5py[version='>=2.9',build=mpi*]
   - netcdf4
-  - pytorch
+  - pytorch>=1.11.0,<2.2.3[channel=pytorch]
   - torchvision
   - scipy
   - pre-commit

scripts/heat_dev_cuda11.yml

@@ -0,0 +1,19 @@
+name: heat_dev_cuda11
+channels:
+  - conda-forge
+  - defaults
+  - pytorch
+  - nvidia
+dependencies:
+  - python=3.11
+  - openmpi
+  - mpi4py
+  - h5py[version='>=2.9',build=mpi*]
+  - netcdf4
+  - pytorch=2.2.2
+  - pytorch-cuda=11.8
+  - torchvision
+  - scipy
+  - pre-commit
+  - black
+  - flake8

scripts/heat_dev_cuda12.yml

@@ -0,0 +1,19 @@
+name: heat_dev_cuda12
+channels:
+  - conda-forge
+  - defaults
+  - pytorch
+  - nvidia
+dependencies:
+  - python=3.11
+  - openmpi
+  - mpi4py
+  - h5py[version='>=2.9',build=mpi*]
+  - netcdf4
+  - pytorch=2.2.2
+  - pytorch-cuda=12.1
+  - torchvision
+  - scipy
+  - pre-commit
+  - black
+  - flake8