diff --git a/.buildkite/bin/setup b/.buildkite/bin/setup index 7d2295db7..aac3ec205 100755 --- a/.buildkite/bin/setup +++ b/.buildkite/bin/setup @@ -8,10 +8,10 @@ export TORCH_EXTENSIONS_DIR=~/$BUILDKITE_BUILD_ID/$BUILDKITE_AGENT_NAME/torch_ex # create env conda create -n $BUILDKITE_BUILD_ID -y -conda env update -n $BUILDKITE_BUILD_ID --file environments/env-dev.yml --prune +conda env update -n $BUILDKITE_BUILD_ID --file environments/linux-cuda/env.yml --prune conda activate $BUILDKITE_BUILD_ID set -x # install pkg -pip install . +pip install .[dev] diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 000000000..51049917c --- /dev/null +++ b/Dockerfile @@ -0,0 +1,24 @@ +# syntax=docker/dockerfile:1 +FROM mambaorg/micromamba:jammy-cuda-12.1.1 + +USER root + +RUN apt-get update && apt-get install -y --no-install-recommends \ + curl \ + unzip \ + git \ + vim \ + build-essential \ + && rm -rf /var/lib/apt/lists/* + +USER $MAMBA_USER + +ENV PATH "$MAMBA_ROOT_PREFIX/bin:$PATH" + +WORKDIR /code + +COPY --chown=$MAMBA_USER:$MAMBA_USER environments/linux-cuda /code/environments/linux-cuda +RUN micromamba install -y -n base -f /code/environments/linux-cuda/env.yml && micromamba clean --all --yes +RUN pip install -r /code/environments/linux-cuda/requirements-dev-linux-cuda.txt +COPY --chown=$MAMBA_USER:$MAMBA_USER . /code +RUN pip install .[dev] diff --git a/environments/devel/Dockerfile b/environments/devel/Dockerfile deleted file mode 100644 index 8b5938f84..000000000 --- a/environments/devel/Dockerfile +++ /dev/null @@ -1,58 +0,0 @@ -# syntax=docker/dockerfile:1 -FROM nvidia/cuda:11.8.0-devel-ubuntu20.04 as base - -# general environment for docker -ENV DEBIAN_FRONTEND=noninteractive - -RUN apt-get update; apt-get install -y --no-install-recommends sudo wget curl git vim rsync libxml2 && rm -rf /var/lib/apt/lists/* - -ENV CONDA_DIR=/opt/conda -ENV PATH ${CONDA_DIR}/bin:$PATH - -RUN curl -LO https://github.com/conda-forge/miniforge/releases/latest/download/Mambaforge-Linux-x86_64.sh \ - && bash Mambaforge-Linux-x86_64.sh -b -p $CONDA_DIR \ - && rm Mambaforge-Linux-x86_64.sh \ - && conda clean -afy \ - && printf "source ${CONDA_DIR}/etc/profile.d/conda.sh\nsource ${CONDA_DIR}/etc/profile.d/mamba.sh\nmamba activate base" >> /etc/skel/.bashrc \ - && printf "source ${CONDA_DIR}/etc/profile.d/conda.sh\nsource ${CONDA_DIR}/etc/profile.d/mamba.sh\nmamba activate base" >> ~/.bashrc - -# create group for conda install -RUN groupadd conda \ - && chgrp -R conda ${CONDA_DIR} \ - && chmod 770 -R ${CONDA_DIR} - -# create docker user -RUN useradd -m -s /bin/bash docker && echo "docker:docker" | chpasswd && adduser docker sudo && adduser docker conda - -# enable passwordless sudo -RUN echo "docker ALL=(ALL) NOPASSWD:ALL" > /etc/sudoers.d/docker - -USER docker -WORKDIR /home/docker - -# FIXME needed for newer versions of sparse but should be replaced by explicit casts to dense where needed -ENV SPARSE_AUTO_DENSIFY 1 - -FROM base as builder - -# NOTE we copy requirements over first so layers are cached and we don't have -# to reinstall dependencies if only source has changed - -### NOTE Dependencies can be isntalled in two different ways - -### 1. install a fresh environment by re-resolving dependencies based on specification in env.yml and requirements{-dev}.in ### - -COPY --chown=docker environments/requirements.in environments/env.yml /home/docker/tmol/ -RUN mamba env update -n base -f /home/docker/tmol/env.yml -COPY --chown=docker environments/requirements-dev.in /home/docker/tmol/ -RUN pip install -r /home/docker/tmol/requirements-dev.in - -### 2. install a frozen environment with exact versions specified from env-dev.yml ### - -# we choose (2) here by default for reproducibility -# COPY --chown=docker env-dev.yml /home/docker/tmol/ -# RUN mamba env update -n base -f /home/docker/tmol/env-dev.yml - -# optional: install tmol into the docker image; assumes that this Dockerfile lives in tmol/environments/test -# COPY --chown=docker ../../. /home/docker/tmol -# RUN pip install -e /home/docker/tmol diff --git a/environments/env-dev.yml b/environments/env-dev.yml deleted file mode 100644 index 1f639f6fb..000000000 --- a/environments/env-dev.yml +++ /dev/null @@ -1,149 +0,0 @@ -name: tmol -channels: - - pytorch - - conda-forge -dependencies: - - _libgcc_mutex=0.1=conda_forge - - _openmp_mutex=4.5=2_kmp_llvm - - blas=1.0=mkl - - brotlipy=0.7.0=py39hb9d737c_1005 - - bzip2=1.0.8=h7f98852_4 - - ca-certificates=2022.12.7=ha878542_0 - - cffi=1.15.1=py39he91dace_3 - - cfgv=3.3.1=pyhd8ed1ab_0 - - clang-format=16.0.2=default_h83cc7fd_0 - - clang-format-16=16.0.2=default_h83cc7fd_0 - - codecov=2.1.13 - - cryptography=40.0.2=py39h079d5ae_0 - - cudatoolkit-dev=11.3.1=py39h3811e60_0 - - distlib=0.3.6=pyhd8ed1ab_0 - - filelock=3.12.0=pyhd8ed1ab_0 - - icu=72.1=hcb278e6_0 - - identify=2.5.23=pyhd8ed1ab_0 - # - intel-openmp=2022.1.0=h9e868ea_3769 - - ld_impl_linux-64=2.40=h41732ed_0 - - libblas=3.9.0=16_linux64_mkl - - libcblas=3.9.0=16_linux64_mkl - - libclang=16.0.2=default_h83cc7fd_0 - - libclang-cpp16=16.0.2=default_h83cc7fd_0 - - libffi=3.4.2=h7f98852_5 - - libgcc-ng=12.2.0=h65d4601_19 - - libgomp=12.2.0=h65d4601_19 - - libiconv=1.17=h166bdaf_0 - - liblapack=3.9.0=16_linux64_mkl - - libllvm10=10.0.1=he513fc3_3 - - libllvm16=16.0.2=hbf9e925_0 - - libnsl=2.0.0=h7f98852_0 - - libsqlite=3.40.0=h753d276_0 - - libstdcxx-ng=12.2.0=h46fd767_19 - - libuuid=2.38.1=h0b41bf4_0 - - libuv=1.44.2=h166bdaf_0 - - libxml2=2.10.4=hfdac1af_0 - - libzlib=1.2.13=h166bdaf_4 - - llvmlite=0.40.0 - - mkl=2022.1.0 - - ncurses=6.3=h27087fc_1 - - ninja=1.11.1=h924138e_0 - - nodeenv=1.7.0=pyhd8ed1ab_0 - - numba=0.57.0 - - numpy=1.24.3 - - openssl=3.1.0=hd590300_2 - - pip=23.1.2=pyhd8ed1ab_0 - - pre-commit=3.2.2 - - pycparser=2.21=pyhd8ed1ab_0 - - pyopenssl=23.1.1=pyhd8ed1ab_0 - - pysocks=1.7.1=pyha2e5f31_6 - - python=3.9.16=h2782a2a_0_cpython - - python_abi=3.9 - - pytorch=1.11.0=py3.9_cuda11.3_cudnn8.2.0_0 - - pytorch-mutex=1.0=cuda - - readline=8.2=h8228510_1 - - setuptools=67.7.2=pyhd8ed1ab_0 - - sysroot_linux-64=2.12=he073ed8_15 - # - tbb=2020.3 - - tk=8.6.12=h27826a3_0 - - typing_extensions=4.5.0=pyha770c72_0 - - ukkonen=1.0.1=py39hf939315_3 - - virtualenv=20.23.0=pyhd8ed1ab_0 - - wheel=0.40.0=pyhd8ed1ab_0 - - xz=5.2.6=h166bdaf_0 - - zlib=1.2.13=h166bdaf_4 - - zstd=1.5.2=h3eb15da_6 - - pip: - - asciitree==0.3.3 - - astor==0.8.1 - - asttokens==2.2.1 - - attrs==23.1.0 - - attrs-strict==1.0.0 - - backcall==0.2.0 - - black==23.3.0 - - cattrs==22.2.0 - - certifi==2022.12.7 - - charset-normalizer==3.1.0 - - click==8.1.3 - - colorama==0.4.6 - - coverage==7.2.4 - - decorator==5.1.1 - - docopt==0.6.2 - - entrypoints==0.4 - - exceptiongroup==1.1.1 - - executing==1.2.0 - - fasteners==0.18 - - flake8==6.0.0 - - frozendict==2.3.7 - - hypothesis==6.70.2 - - idna==3.4 - - itermplot==0.5 - - iniconfig==2.0.0 - - ipython==8.13.0 - - jedi==0.18.2 - # - matplotlib==3.7.1 - # - matplotlib-inline==0.1.6 - - mccabe==0.7.0 - - mypy-extensions==1.0.0 - - networkx==3.1 - - numcodecs==0.11.0 - - packaging==23.1 - - pandas==2.0.1 - - parso==0.8.3 - - pathspec==0.11.1 - - pexpect==4.8.0 - - pickleshare==0.7.5 - - pint==0.20.1 - - platformdirs==3.5.0 - - pluggy==1.0.0 - - prompt-toolkit==3.0.38 - - psutil==5.9.5 - - ptyprocess==0.7.0 - - pure-eval==0.2.2 - - py==1.11.0 - - py-cpuinfo==9.0.0 - - pycodestyle==2.10.0 - - pyflakes==3.0.1 - - pygments==2.15.1 - - pytest==7.3.1 - - pytest-benchmark==4.0.0 - - pytest-cov==4.0.0 - - pytest-forked==1.6.0 - - pytest-instafail==0.5.0 - - pytest-repeat==0.9.1 - - pytest-watch==4.2.0 - - python-dateutil==2.8.2 - - pytz==2023.3 - - pyyaml==6.0 - - requests==2.28.2 - - scipy==1.10.1 - - seaborn==0.12.2 - - six==1.16.0 - - sortedcontainers==2.4.0 - - stack-data==0.6.2 - - tomli==2.0.1 - - toolz==0.12.0 - - traitlets==5.9.0 - - typing-inspect==0.8.0 - - typish==1.9.3 - - tzdata==2023.3 - - urllib3==1.26.15 - - watchdog==3.0.0 - - wcwidth==0.2.6 - - zarr==2.14.2 diff --git a/environments/env.yml b/environments/env.yml deleted file mode 100644 index cf9f873f7..000000000 --- a/environments/env.yml +++ /dev/null @@ -1,17 +0,0 @@ -name: tmol -channels: - - pytorch - - conda-forge -dependencies: - - python=3.9 - - pytorch==1.11.0 - - cudatoolkit-dev=11.3 - - clang-format - - codecov - - numpy - - numba - - ninja - - pip - - pre-commit - - pip: - - -r requirements.in diff --git a/environments/linux-cuda/build_requirements.sh b/environments/linux-cuda/build_requirements.sh new file mode 100755 index 000000000..153fc9aaa --- /dev/null +++ b/environments/linux-cuda/build_requirements.sh @@ -0,0 +1,4 @@ +#!/usr/bin/env bash + +pip-compile --verbose --no-emit-index-url --resolver=backtracking -o requirements-linux-cuda.txt ../../requirements.in +pip-compile --verbose --no-emit-index-url --resolver=backtracking -c requirements-linux-cuda.txt -o requirements-dev-linux-cuda.txt ../../requirements-dev.in diff --git a/environments/linux-cuda/env.yml b/environments/linux-cuda/env.yml new file mode 100644 index 000000000..7f659f05a --- /dev/null +++ b/environments/linux-cuda/env.yml @@ -0,0 +1,10 @@ +name: tmol +channels: + - nvidia/label/cuda-12.1.1 + - conda-forge +dependencies: + - python=3.11 + - cuda + - pip + - pip: + - -r requirements-linux-cuda.txt diff --git a/environments/linux-cuda/requirements-dev-linux-cuda.txt b/environments/linux-cuda/requirements-dev-linux-cuda.txt new file mode 100644 index 000000000..71ebcd23e --- /dev/null +++ b/environments/linux-cuda/requirements-dev-linux-cuda.txt @@ -0,0 +1,159 @@ +# +# This file is autogenerated by pip-compile with Python 3.11 +# by the following command: +# +# pip-compile --constraint=requirements-linux-cuda.txt --no-emit-index-url --output-file=requirements-dev-linux-cuda.txt ../../requirements-dev.in +# +--trusted-host data.pyg.org + +black==24.1.1 + # via -r ../../requirements-dev.in +build==1.0.3 + # via pip-tools +certifi==2024.2.2 + # via + # -c requirements-linux-cuda.txt + # requests +charset-normalizer==3.3.2 + # via + # -c requirements-linux-cuda.txt + # requests +click==8.1.7 + # via + # black + # pip-tools +codecov==2.1.13 + # via -r ../../requirements-dev.in +colorama==0.4.6 + # via pytest-watch +contourpy==1.2.0 + # via matplotlib +coverage[toml]==7.4.1 + # via + # codecov + # pytest-cov +cycler==0.12.1 + # via matplotlib +docopt==0.6.2 + # via pytest-watch +flake8==7.0.0 + # via -r ../../requirements-dev.in +fonttools==4.48.1 + # via matplotlib +idna==3.6 + # via + # -c requirements-linux-cuda.txt + # requests +iniconfig==2.0.0 + # via pytest +itermplot==0.5 + # via -r ../../requirements-dev.in +kiwisolver==1.4.5 + # via matplotlib +matplotlib==3.8.2 + # via + # itermplot + # seaborn +mccabe==0.7.0 + # via flake8 +mypy-extensions==1.0.0 + # via + # -c requirements-linux-cuda.txt + # black +numpy==1.26.4 + # via + # -c requirements-linux-cuda.txt + # contourpy + # itermplot + # matplotlib + # pandas + # seaborn +packaging==23.2 + # via + # black + # build + # matplotlib + # pytest +pandas==2.2.0 + # via + # -c requirements-linux-cuda.txt + # seaborn +pathspec==0.12.1 + # via black +pillow==10.2.0 + # via matplotlib +pip-tools==7.3.0 + # via -r ../../requirements-dev.in +platformdirs==4.2.0 + # via black +pluggy==1.4.0 + # via pytest +py==1.11.0 + # via pytest-forked +py-cpuinfo==9.0.0 + # via pytest-benchmark +pycodestyle==2.11.1 + # via flake8 +pyflakes==3.2.0 + # via flake8 +pyparsing==3.1.1 + # via matplotlib +pyproject-hooks==1.0.0 + # via build +pytest==8.0.0 + # via + # -r ../../requirements-dev.in + # pytest-benchmark + # pytest-cov + # pytest-forked + # pytest-instafail + # pytest-repeat + # pytest-watch +pytest-benchmark==4.0.0 + # via -r ../../requirements-dev.in +pytest-cov==4.1.0 + # via -r ../../requirements-dev.in +pytest-forked==1.6.0 + # via -r ../../requirements-dev.in +pytest-instafail==0.5.0 + # via -r ../../requirements-dev.in +pytest-repeat==0.9.3 + # via -r ../../requirements-dev.in +pytest-watch==4.2.0 + # via -r ../../requirements-dev.in +python-dateutil==2.8.2 + # via + # -c requirements-linux-cuda.txt + # matplotlib + # pandas +pytz==2024.1 + # via + # -c requirements-linux-cuda.txt + # pandas +requests==2.31.0 + # via + # -c requirements-linux-cuda.txt + # codecov +seaborn==0.13.2 + # via -r ../../requirements-dev.in +six==1.16.0 + # via + # -c requirements-linux-cuda.txt + # itermplot + # python-dateutil +tzdata==2023.4 + # via + # -c requirements-linux-cuda.txt + # pandas +urllib3==2.2.0 + # via + # -c requirements-linux-cuda.txt + # requests +watchdog==4.0.0 + # via pytest-watch +wheel==0.42.0 + # via pip-tools + +# The following packages are considered to be unsafe in a requirements file: +# pip +# setuptools diff --git a/environments/linux-cuda/requirements-linux-cuda.txt b/environments/linux-cuda/requirements-linux-cuda.txt new file mode 100644 index 000000000..224b6ad11 --- /dev/null +++ b/environments/linux-cuda/requirements-linux-cuda.txt @@ -0,0 +1,185 @@ +# +# This file is autogenerated by pip-compile with Python 3.11 +# by the following command: +# +# pip-compile --no-emit-index-url --output-file=requirements-linux-cuda.txt ../../requirements.in +# +--trusted-host data.pyg.org + +asciitree==0.3.3 + # via zarr +astor==0.8.1 + # via -r ../../requirements.in +asttokens==2.4.1 + # via stack-data +attrs==23.2.0 + # via + # -r ../../requirements.in + # attrs-strict + # cattrs + # hypothesis +attrs-strict==1.0.1 + # via -r ../../requirements.in +cattrs==23.2.3 + # via -r ../../requirements.in +certifi==2024.2.2 + # via requests +charset-normalizer==3.3.2 + # via requests +decorator==5.1.1 + # via + # -r ../../requirements.in + # ipython +executing==2.0.1 + # via stack-data +fasteners==0.19 + # via zarr +filelock==3.13.1 + # via + # torch + # triton +frozendict==2.4.0 + # via -r ../../requirements.in +fsspec==2024.2.0 + # via torch +hypothesis==6.98.3 + # via -r ../../requirements.in +idna==3.6 + # via requests +ipython==8.21.0 + # via -r ../../requirements.in +jedi==0.19.1 + # via ipython +jinja2==3.1.3 + # via torch +llvmlite==0.42.0 + # via numba +markupsafe==2.1.5 + # via jinja2 +matplotlib-inline==0.1.6 + # via ipython +mpmath==1.3.0 + # via sympy +mypy-extensions==1.0.0 + # via typing-inspect +networkx==3.2.1 + # via + # -r ../../requirements.in + # torch +ninja==1.11.1.1 + # via -r ../../requirements.in +numba==0.59.0 + # via sparse +numcodecs==0.12.1 + # via zarr +numpy==1.26.4 + # via + # -r ../../requirements.in + # numba + # numcodecs + # pandas + # pyarrow + # scipy + # sparse + # zarr +nvidia-cublas-cu12==12.1.3.1 + # via + # nvidia-cudnn-cu12 + # nvidia-cusolver-cu12 + # torch +nvidia-cuda-cupti-cu12==12.1.105 + # via torch +nvidia-cuda-nvrtc-cu12==12.1.105 + # via torch +nvidia-cuda-runtime-cu12==12.1.105 + # via torch +nvidia-cudnn-cu12==8.9.2.26 + # via torch +nvidia-cufft-cu12==11.0.2.54 + # via torch +nvidia-curand-cu12==10.3.2.106 + # via torch +nvidia-cusolver-cu12==11.4.5.107 + # via torch +nvidia-cusparse-cu12==12.1.0.106 + # via + # nvidia-cusolver-cu12 + # torch +nvidia-nccl-cu12==2.19.3 + # via torch +nvidia-nvjitlink-cu12==12.3.101 + # via + # nvidia-cusolver-cu12 + # nvidia-cusparse-cu12 +nvidia-nvtx-cu12==12.1.105 + # via torch +pandas==2.2.0 + # via -r ../../requirements.in +parso==0.8.3 + # via jedi +pexpect==4.9.0 + # via ipython +pint==0.23 + # via -r ../../requirements.in +prompt-toolkit==3.0.43 + # via ipython +psutil==5.9.8 + # via -r ../../requirements.in +ptyprocess==0.7.0 + # via pexpect +pure-eval==0.2.2 + # via stack-data +pyarrow==15.0.0 + # via -r ../../requirements.in +pygments==2.17.2 + # via ipython +python-dateutil==2.8.2 + # via pandas +pytz==2024.1 + # via pandas +pyyaml==6.0.1 + # via -r ../../requirements.in +requests==2.31.0 + # via -r ../../requirements.in +scipy==1.12.0 + # via + # -r ../../requirements.in + # sparse +six==1.16.0 + # via + # asttokens + # python-dateutil +sortedcontainers==2.4.0 + # via hypothesis +sparse==0.15.1 + # via -r ../../requirements.in +stack-data==0.6.3 + # via ipython +sympy==1.12 + # via torch +toolz==0.12.1 + # via -r ../../requirements.in +torch==2.2.0 + # via -r ../../requirements.in +traitlets==5.14.1 + # via + # ipython + # matplotlib-inline +triton==2.2.0 + # via torch +typing-extensions==4.9.0 + # via + # -r ../../requirements.in + # pint + # torch + # typing-inspect +typing-inspect==0.9.0 + # via -r ../../requirements.in +tzdata==2023.4 + # via pandas +urllib3==2.2.0 + # via requests +wcwidth==0.2.13 + # via prompt-toolkit +zarr==2.16.1 + # via -r ../../requirements.in diff --git a/environments/test/Dockerfile b/environments/test/Dockerfile deleted file mode 100644 index de56e603d..000000000 --- a/environments/test/Dockerfile +++ /dev/null @@ -1,81 +0,0 @@ -# syntax=docker/dockerfile:1 -FROM nvidia/cuda:11.8.0-devel-ubuntu20.04 as base -ENV LANG=C.UTF-8 LC_ALL=C.UTF-8 - -# general environment for docker -ENV DEBIAN_FRONTEND=noninteractive - -RUN apt-get update --fix-missing && \ - apt-get install -y --no-install-recommends sudo wget curl bzip2 ca-certificates git vim rsync jq software-properties-common && \ - apt-get clean && \ - rm -rf /var/lib/apt/lists/* - -ENV CONDA_DIR=/opt/conda -ENV PATH ${CONDA_DIR}/bin:$PATH - -RUN curl -LO https://github.com/conda-forge/miniforge/releases/latest/download/Mambaforge-Linux-x86_64.sh \ - && bash Mambaforge-Linux-x86_64.sh -b -p $CONDA_DIR \ - && rm Mambaforge-Linux-x86_64.sh \ - && conda clean -afy \ - && printf "source ${CONDA_DIR}/etc/profile.d/conda.sh\nsource ${CONDA_DIR}/etc/profile.d/mamba.sh\nmamba activate base" >> /etc/skel/.bashrc \ - && printf "source ${CONDA_DIR}/etc/profile.d/conda.sh\nsource ${CONDA_DIR}/etc/profile.d/mamba.sh\nmamba activate base" >> ~/.bashrc - -# create group for conda install -RUN groupadd conda \ - && chgrp -R conda ${CONDA_DIR} \ - && chmod 770 -R ${CONDA_DIR} - -# create docker user -RUN useradd -m -s /bin/bash docker && echo "docker:docker" | chpasswd && adduser docker sudo && adduser docker conda - -# enable passwordless sudo -RUN echo "docker ALL=(ALL) NOPASSWD:ALL" > /etc/sudoers.d/docker - - -# Inlined tini version in fetch url -# Question: why are we writing directly to /usr/bin? -ENV TINI_VERSION v0.16.1 -RUN wget --quiet -O /usr/bin/tini https://github.com/krallin/tini/releases/download/v0.16.1/tini -RUN chmod +x /usr/bin/tini - -# create the entry point now before we lose sudo strength -RUN echo "#!/bin/bash" > /docker-entrypoint.sh && \ - echo ". /opt/conda/etc/profile.d/conda.sh" >> /docker-entrypoint.sh && \ - echo "conda activate tmol" >> /docker-entrypoint.sh && \ - echo "exec \"\$@\"" >> /docker-entrypoint.sh && \ - chmod a+rx /docker-entrypoint.sh - -#USER docker -#WORKDIR /home/docker - -# FIXME needed for newer versions of sparse but should be replaced by explicit casts to dense where needed -ENV SPARSE_AUTO_DENSIFY 1 - -FROM base as builder - - -# NOTE we copy requirements over first so layers are cached and we don't have -# to reinstall dependencies if only source has changed - -### NOTE Dependencies can be installed in two different ways - -### 1. install a fresh environment by re-resolving dependencies based on specification in env.yml and requirements{-dev}.in ### - -# COPY --chown=docker requirements.in env.yml /home/docker/tmol/ -# RUN mamba env update -n base -f /home/docker/tmol/env.yml -# COPY --chown=docker requirements-dev.in /home/docker/tmol/ -# RUN pip install -r /home/docker/tmol/requirements-dev.in - -### 2. install a frozen environment with exact versions specified from env-dev.yml ### -# For the testing server, we choose (2) here for reproducibility - -#COPY --chown=docker env-dev.yml /home/docker/tmol/ -COPY env-dev.yml /home/docker/tmol/ -RUN mamba create -n tmol && \ - mamba env update -n tmol -f /home/docker/tmol/env-dev.yml && \ - mamba clean -afy - - -ENTRYPOINT [ "/usr/bin/tini", "--", "/docker-entrypoint.sh" ] -CMD [ "/bin/bash" ] - diff --git a/environments/test/env-dev.yml b/environments/test/env-dev.yml deleted file mode 120000 index cd5309890..000000000 --- a/environments/test/env-dev.yml +++ /dev/null @@ -1 +0,0 @@ -../env-dev.yml \ No newline at end of file diff --git a/environments/test/requirements-dev.in b/environments/test/requirements-dev.in deleted file mode 120000 index a92438b65..000000000 --- a/environments/test/requirements-dev.in +++ /dev/null @@ -1 +0,0 @@ -../requirements-dev.in \ No newline at end of file diff --git a/environments/test/requirements.in b/environments/test/requirements.in deleted file mode 120000 index 974482ec2..000000000 --- a/environments/test/requirements.in +++ /dev/null @@ -1 +0,0 @@ -../requirements.in \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index bfaddfe62..5a31fea4b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,21 +1,30 @@ -[tool.black] -py36 = true -include = '\.pyi?$' -exclude = ''' -/( - \.git - | \.hg - | \.mypy_cache - | \.tox - | \.venv - | \.conda - | _build - | buck-out - | build - | dist - # Project specific - | tmol/extern - | dev -)/ -''' +[project] +name = "tmol" +dynamic = ["version", "readme", "dependencies", "optional-dependencies"] +requires-python = ">=3.9" +[tool.setuptools.dynamic] +dependencies = {file = ["requirements.in"]} +readme = {file = "README.md"} +optional-dependencies = {dev = {file = ["requirements-dev.in"]}} + +[tool.setuptools] +packages = ["tmol"] + +[tool.setuptools.package-data] +"*" = ["*.hh", "*.cpp", "*.cu", "*.cuh", "*.cc", ".hpp", "*.h", "*.hxx"] +"tmol.database.default.chemical" = ["**/*.yaml"] +"tmol.database.default.scoring" = ["**/*.yaml"] +"tmol.tests.data.pdb" = ["**/*.pdb"] +"tmol.tests.data.rosetta_baselines" = ["**/*.pickle"] +"tmol.tests.data.constraints" = ["**/*.fa", "**/*.npz"] + +[build-system] +requires = ["setuptools >= 65", "setuptools_scm[toml]>=6.2"] +build-backend = 'setuptools.build_meta' + +[tool.setuptools_scm] +search_parent_directories = true +version_scheme = "no-guess-dev" +local_scheme = "node-and-date" +fallback_version = "0.0.0" diff --git a/environments/requirements-dev.in b/requirements-dev.in similarity index 82% rename from environments/requirements-dev.in rename to requirements-dev.in index 8f80b4c71..7100f77e8 100644 --- a/environments/requirements-dev.in +++ b/requirements-dev.in @@ -1,4 +1,5 @@ black +codecov flake8 itermplot==0.5 pytest @@ -8,4 +9,5 @@ pytest-instafail pytest-repeat pytest-watch pytest-benchmark -seaborn \ No newline at end of file +seaborn +pip-tools diff --git a/environments/requirements.in b/requirements.in similarity index 82% rename from environments/requirements.in rename to requirements.in index af17013fb..91036928a 100644 --- a/environments/requirements.in +++ b/requirements.in @@ -1,4 +1,5 @@ -#numpy # NOTE we get numpy from conda to ensure compat with numba +torch +numpy # NOTE we get numpy from conda to ensure compat with numba astor attrs != 22.2.0 # problem unpickling tmol/tests/data/rosetta_baseline/1ubq.scores.pickle in 22.2.0 (see https://github.com/python-attrs/attrs/pull/1085) attrs_strict @@ -9,6 +10,7 @@ hypothesis ipython networkx pandas +pyarrow pint psutil pyyaml @@ -18,5 +20,5 @@ sparse # NOTE sparse>0.3.1 requires env var SPARSE_AUTO_DENSIFY=1 to be set toolz typing_extensions typing_inspect -typish zarr +ninja diff --git a/setup.py b/setup.py deleted file mode 100755 index a21c4aeb5..000000000 --- a/setup.py +++ /dev/null @@ -1,73 +0,0 @@ -#!/usr/bin/env python - -import sys -from setuptools import setup, find_packages -import subprocess -import re -import os - - -def git_version(): - try: - git_describe = ( - subprocess.check_output( - ["git", "describe", "--long", "--tags", "--match", "[0-9]*"] - ) - .strip() - .decode() - ) - except subprocess.CalledProcessError: - version = "0.0.0" - return version - - describe_match = re.match( - r"(?P[0-9.]+)(-(?P\d+)-g(?P\w+))?", git_describe - ) - if not describe_match: - raise ValueError("Invalid version.", git_describe) - else: - desc = describe_match.groupdict() - - desc["post_revision"] = int(desc.get("post_revision", 0)) - if not desc["post_revision"]: - version = f"{desc['version']}+{desc['commit']}" - else: - version = f"{desc['version']}.post.dev+{desc['post_revision']}.{desc['commit']}" - - return version - - -def find_cpp_files(directory): - paths = [] - for path, directories, filenames in os.walk(directory): - for filename in filenames: - _, ext = os.path.splitext(filename) - if ext in [".hh", ".cpp", ".cu", ".cuh", ".cc", ".hpp", ".h", ".hxx"]: - paths.append(os.path.join("..", path, filename)) - return paths - - -extra_files = find_cpp_files(".") -extra_files.extend( - [ - "../tmol/database/default/chemical/*", - "../tmol/database/default/scoring/*", - "../tmol/tests/data/pdb/*", - "../tmol/tests/data/pdb/*", - "../tmol/tests/data/rosetta_baseline/*", - "../tmol/tests/data/constraints/*", - ] -) - - -needs_pytest = {"pytest", "test"}.intersection(sys.argv) -pytest_runner = ["pytest-runner"] if needs_pytest else [] - -setup( - name="tmol", - version=git_version(), - packages=find_packages(), - package_data={"": extra_files}, - setup_requires=pytest_runner, - zip_safe=False, -) diff --git a/tmol/__init__.py b/tmol/__init__.py index ffe465e4e..f49dff21e 100644 --- a/tmol/__init__.py +++ b/tmol/__init__.py @@ -1,3 +1,5 @@ +from importlib.metadata import PackageNotFoundError, version + from tmol.chemical.restypes import three2one # noqa: F401 from tmol.io import pose_stack_from_pdb # noqa: F401 from tmol.io.pose_stack_construction import ( # noqa: F401 @@ -16,6 +18,7 @@ ) from tmol.io.pose_stack_from_rosettafold2 import ( # noqa: F401 pose_stack_from_rosettafold2, + pose_stack_to_rosettafold2, canonical_form_from_rosettafold2, canonical_ordering_for_rosettafold2, packed_block_types_for_rosettafold2, @@ -24,9 +27,21 @@ write_pose_stack_pdb, atom_records_from_pose_stack, # TO DO: SHOULD THIS BE IN THE API?? ) + +from tmol.optimization.sfxn_modules import CartesianSfxnNetwork as cart_sfxn_network + +from tmol.optimization.lbfgs_armijo import LBFGS_Armijo as lbfgs_armijo + + from tmol.score import beta2016_score_function # noqa: F401 +try: + __version__ = version("tmol") +except PackageNotFoundError: + __version__ = "unknown version" + + def include_paths(): """C++/CUDA include paths for tmol components.""" diff --git a/tmol/io/pose_stack_from_rosettafold2.py b/tmol/io/pose_stack_from_rosettafold2.py index 301a25850..78334a9d7 100644 --- a/tmol/io/pose_stack_from_rosettafold2.py +++ b/tmol/io/pose_stack_from_rosettafold2.py @@ -130,6 +130,80 @@ def canonical_form_from_rosettafold2(seq, xyz, chainlens): ) +def pose_stack_to_rosettafold2(seq, xyz, chainlens, pose_stack): + """The canonical form is intended to represent a stable, serializable intermediate format + for a structure so that it can be created today and then be read in years from now + and be used to construct a PoseStack in tmol. As residue types are integers, + this means that we must guarantee stability of these integer representations, but it also + means that you the user must build a PoseStack using the carefully constructed objects + returned by the canonical_ordering_for_rosettafold2 and packed_block_types_for_rosettafold2 + functions. + + E.g.: + seq, xyz, chainlens = rosettafold2_model.infer(sequence) + cf = tmol.canonical_form_from_rosettafold2(seq, xyz, chainlens) + torch.save(cf, "saved_canonical_form.pt") + + # then later + cf2 = {x: y.to(device) for x,y in torch.load("saved_canonical_form.pt")} + co = canonical_ordering_for_rosettafold2() + pbt = packed_block_types_for_rosettafold2(device) + pose_stack = tmol.pose_stack_from_canonical_form(co, pbt, **cf2) + + """ + + from tmol.io.pose_stack_deconstruction import canonical_form_from_pose_stack + + device = xyz.device + n_poses = 1 # RF2 does not presently do batch processing + max_n_res = seq.shape[0] + max_n_ats = xyz.shape[1] + + rf2_pose_ind_for_atom = ( + torch.arange(n_poses, dtype=torch.int64, device=device) + .reshape(-1, 1, 1) + .expand(-1, max_n_res, max_n_ats) + ) + rf2_res_ind_for_atom = ( + torch.arange(max_n_res, dtype=torch.int64, device=device) + .reshape(1, -1, 1) + .expand(n_poses, -1, max_n_ats) + ) + + assert device == seq.device + + co = canonical_ordering_for_rosettafold2() + ( + rf22t_rtmap, + rf22t_atmap, + rf2_at_is_real_map, + supress_atom_for_nterm, + ) = _get_rf2_2_tmol_mappings(device) + + canonical_form = canonical_form_from_pose_stack(co, pose_stack) + + seq = seq.unsqueeze(0) + xyz = xyz.unsqueeze(0) + tmol_restypes = rf22t_rtmap[seq] # reverse this mapping + atom_mapping = rf22t_atmap[seq] + rf2_at_is_real = rf2_at_is_real_map[seq] + + rf2_coords = torch.full( # allocate xyz for RF2 instead, with correct size - length (sum(L_s)), 27, 3 + xyz.shape, + numpy.NaN, + dtype=torch.float32, + device=device, + ) + + rf2_coords[rf2_at_is_real] = canonical_form[2][ + rf2_pose_ind_for_atom[rf2_at_is_real], + rf2_res_ind_for_atom[rf2_at_is_real], + atom_mapping[rf2_at_is_real], + ] + + return rf2_coords + + @toolz.functoolz.memoize def _paramdb_for_rosettafold2() -> ParameterDatabase: """Construct the paramdb representing the subset of residues that diff --git a/tmol/optimization/modules.py b/tmol/optimization/modules.py index 218152d86..8245e52ce 100755 --- a/tmol/optimization/modules.py +++ b/tmol/optimization/modules.py @@ -2,7 +2,8 @@ import attr from tmol.system.kinematics import KinematicDescription -from tmol.system.score_support import kincoords_to_coords + +# from tmol.system.score_support import kincoords_to_coords # causes circular import when importing tmol from RF2 # modules for cartesian and torsion-space optimization # @@ -74,6 +75,7 @@ def torsional_energy_network_from_system( ) +""" # torsion space minimization class TorsionalEnergyNetwork(torch.nn.Module): def __init__(self, score_system, dofs, kinforest, system_size, dof_mask=None): @@ -103,3 +105,4 @@ def coords(self): def forward(self): return self.score_system.intra_total(self.coords()) +""" diff --git a/tmol/pack/rotamer/build_rotamers.py b/tmol/pack/rotamer/build_rotamers.py index eb2a46771..e1fb5a077 100644 --- a/tmol/pack/rotamer/build_rotamers.py +++ b/tmol/pack/rotamer/build_rotamers.py @@ -524,7 +524,7 @@ def merge_chi_samples(chi_samples): ) sampler_for_rotamer_unsorted = torch.cat( [ - torch.full((samples[1].shape[0],), i, dtype=torch.int64) + torch.full((samples[1].shape[0],), i, dtype=torch.int64, device=device) for i, samples in enumerate(chi_samples) ] ) @@ -668,7 +668,7 @@ def create_dof_inds_to_copy_from_orig_to_rotamers( rot_mcfp_at_inds_kto[rot_mcfp_at_inds_kto != -1] += n_dof_atoms_offset_for_rot[ torch.div( - torch.arange(n_rots * max_n_mcfp_atoms, dtype=torch.int64), + torch.arange(n_rots * max_n_mcfp_atoms, dtype=torch.int64, device=poses.device), max_n_mcfp_atoms, rounding_mode="trunc", )[rot_mcfp_at_inds_kto != -1] diff --git a/tmol/pack/rotamer/dunbrack/dunbrack_chi_sampler.py b/tmol/pack/rotamer/dunbrack/dunbrack_chi_sampler.py index b623511df..82f59528a 100644 --- a/tmol/pack/rotamer/dunbrack/dunbrack_chi_sampler.py +++ b/tmol/pack/rotamer/dunbrack/dunbrack_chi_sampler.py @@ -627,6 +627,7 @@ def package_samples_for_output( for rt in rlt.allowed_restypes ], dtype=torch.uint8, + device=self.device, ) n_restypes_total = restype_is_allowed_for_dun.shape[0] dun_allowed_inds = torch.nonzero(restype_is_allowed_for_dun)[:, 0] diff --git a/tmol/score/common/stack_condense.py b/tmol/score/common/stack_condense.py index 09fde24b7..bee8ad444 100644 --- a/tmol/score/common/stack_condense.py +++ b/tmol/score/common/stack_condense.py @@ -401,7 +401,7 @@ def _value_or_arg_tile_subset_indices( return_args: bool, max_entry: Optional[int] = None, ): - if type(indices) == torch.Tensor: + if type(indices) is torch.Tensor: if max_entry is None: max_entry = torch.max(indices) n_tiles = max_entry // tile_size + 1 @@ -415,7 +415,7 @@ def _value_or_arg_tile_subset_indices( ind_arange = torch.arange( indices.shape[0], dtype=indices.dtype, device=indices.device ) - elif type(indices) == numpy.ndarray: + elif type(indices) is numpy.ndarray: if max_entry is None: max_entry = numpy.amax(indices) n_tiles = max_entry // tile_size + 1 @@ -427,7 +427,7 @@ def _value_or_arg_tile_subset_indices( raise ValueError for i in range(n_tiles): subset = (indices >= i * tile_size) & (indices < (i + 1) * tile_size) - if type(tiled_indices) == torch.Tensor: + if type(tiled_indices) is torch.Tensor: subset_size = torch.sum(subset).cpu() else: subset_size = numpy.sum(subset) diff --git a/tmol/score/ljlk/potentials/params.hh b/tmol/score/ljlk/potentials/params.hh index 446d58da7..da1bd9ad9 100644 --- a/tmol/score/ljlk/potentials/params.hh +++ b/tmol/score/ljlk/potentials/params.hh @@ -170,8 +170,9 @@ struct enable_tensor_view> { template struct enable_tensor_view> { static const bool enabled = enable_tensor_view::enabled; - static const at::ScalarType scalar_type = - enable_tensor_view::scalar_type; + static const at::ScalarType scalar_type() { + return enable_tensor_view::scalar_type(); + } static const int nconsumed_dims = 1; static const int consumed_dims(int i) { diff --git a/tmol/tests/kinematics/segscan/test_segscan.py b/tmol/tests/kinematics/segscan/test_segscan.py index c035d60ee..10cf33b43 100644 --- a/tmol/tests/kinematics/segscan/test_segscan.py +++ b/tmol/tests/kinematics/segscan/test_segscan.py @@ -31,7 +31,7 @@ def test_segscan_inclusive(extension): y = extension.segscan_incl(xcuda, segscuda) y = y.to(device="cpu") - torch.testing.assert_allclose(gold, y) + torch.testing.assert_close(gold, y) @requires_cuda @@ -52,7 +52,7 @@ def test_segscan_exclusive(extension): y = extension.segscan_excl(xcuda, segscuda) y = y.to(device="cpu") - torch.testing.assert_allclose(gold, y) + torch.testing.assert_close(gold, y) @requires_cuda @@ -69,12 +69,12 @@ def test_segscan(extension): y = extension.segscan_incl(x, segs) ycuda = extension.segscan_incl(xcuda, segscuda) - torch.testing.assert_allclose(ycuda.to(device="cpu"), y) + torch.testing.assert_close(ycuda.to(device="cpu"), y) y = extension.segscan_excl(x, segs) ycuda = extension.segscan_excl(xcuda, segscuda) - torch.testing.assert_allclose(ycuda.to(device="cpu"), y) + torch.testing.assert_close(ycuda.to(device="cpu"), y) @requires_cuda @@ -91,12 +91,12 @@ def test_segscan2(extension): y = extension.weird_segscan_incl_128_2(x, segs) ycuda = extension.weird_segscan_incl_128_2(xcuda, segscuda) - torch.testing.assert_allclose(ycuda.to(device="cpu"), y) + torch.testing.assert_close(ycuda.to(device="cpu"), y) y = extension.segscan_excl(x, segs) ycuda = extension.segscan_excl(xcuda, segscuda) - torch.testing.assert_allclose(ycuda.to(device="cpu"), y) + torch.testing.assert_close(ycuda.to(device="cpu"), y) @requires_cuda @@ -498,7 +498,7 @@ def test_segscan_highly_segmented(extension): segscuda = segs.to(device="cuda") y = extension.segscan_incl(x, segs) ycuda = extension.segscan_incl(xcuda, segscuda) - torch.testing.assert_allclose(ycuda.to(device="cpu"), y) + torch.testing.assert_close(ycuda.to(device="cpu"), y) x = torch.ones(840, dtype=torch.float32) segs = torch.tensor( @@ -813,7 +813,7 @@ def test_segscan_highly_segmented(extension): segscuda = segs.to(device="cuda") y = extension.segscan_incl_128_2(x, segs) ycuda = extension.segscan_incl_128_2(xcuda, segscuda) - torch.testing.assert_allclose(ycuda.to(device="cpu"), y) + torch.testing.assert_close(ycuda.to(device="cpu"), y) @requires_cuda @@ -869,7 +869,7 @@ def test_segscan_highly_segmented2(extension): # print(y) # print("y cuda") # print(ycuda.cpu()) - torch.testing.assert_allclose(ycuda.to(device="cpu"), y) + torch.testing.assert_close(ycuda.to(device="cpu"), y) def highly_segmented3(): @@ -5213,7 +5213,7 @@ def test_segscan_highly_segmented3(extension): # print(y) # print("y cuda") # print(ycuda.cpu()) - torch.testing.assert_allclose(ycuda.to(device="cpu"), y) + torch.testing.assert_close(ycuda.to(device="cpu"), y) @requires_cuda diff --git a/tmol/tests/kinematics/test_dof_modules.py b/tmol/tests/kinematics/test_dof_modules.py index a0c954658..deca12edc 100644 --- a/tmol/tests/kinematics/test_dof_modules.py +++ b/tmol/tests/kinematics/test_dof_modules.py @@ -20,8 +20,16 @@ def test_cartesian_coord_factory(ubq_system): # Coords are returned from forward assert src.coords.shape == (1, ubq_system.system_size, 3) - torch.testing.assert_allclose(src.coords[0], ubq_system.coords) - torch.testing.assert_allclose(src()[0], ubq_system.coords) + torch.testing.assert_close( + src.coords[0], + torch.tensor(ubq_system.coords, dtype=src.coords.dtype), + equal_nan=True, + ) + torch.testing.assert_close( + src()[0], + torch.tensor(ubq_system.coords, dtype=src.coords.dtype), + equal_nan=True, + ) # Device defaults and device clone clone = CartesianDOFs.build_from(src) @@ -31,12 +39,12 @@ def test_cartesian_coord_factory(ubq_system): # Coords are copied, not referenced with torch.no_grad(): - torch.testing.assert_allclose(src.coords, clone.coords) + torch.testing.assert_close(src.coords, clone.coords, equal_nan=True) clone.coords[0] += 1 with pytest.raises(AssertionError): - torch.testing.assert_allclose(src.coords, clone.coords) + torch.testing.assert_close(src.coords, clone.coords, equal_nan=True) clone.coords[0] -= 1 - torch.testing.assert_allclose(src.coords, clone.coords) + torch.testing.assert_close(src.coords, clone.coords, equal_nan=True) # Device can be overridden clone = clone.to(cuda_device) @@ -46,8 +54,16 @@ def test_cartesian_coord_factory(ubq_system): assert clone().device == cuda_device # Coords are returned from forward - torch.testing.assert_allclose(clone.coords.cpu()[0], ubq_system.coords) - torch.testing.assert_allclose(clone().cpu()[0], ubq_system.coords) + torch.testing.assert_close( + clone.coords.cpu()[0], + torch.tensor(ubq_system.coords, dtype=clone.coords.dtype), + equal_nan=True, + ) + torch.testing.assert_close( + clone().cpu()[0], + torch.tensor(ubq_system.coords, dtype=clone.coords.dtype), + equal_nan=True, + ) @requires_cuda @@ -57,7 +73,9 @@ def test_kinematic_dof_factory(ubq_system): src = KinematicDOFs.build_from(ubq_system) - torch.testing.assert_allclose(src()[0], ubq_system.coords) + torch.testing.assert_close( + src()[0], torch.tensor(ubq_system.coords), equal_nan=True + ) # Device defaults and device clone clone: KinematicDOFs = KinematicDOFs.build_from(src) @@ -66,16 +84,24 @@ def test_kinematic_dof_factory(ubq_system): # dofs are copied, not referenced with torch.no_grad(): - torch.testing.assert_allclose(src.dofs, clone.dofs) + torch.testing.assert_close(src.dofs, clone.dofs) clone.dofs[0] += 1 with pytest.raises(AssertionError): - torch.testing.assert_allclose(src.dofs, clone.dofs) + torch.testing.assert_close(src.dofs, clone.dofs) with pytest.raises(AssertionError): - torch.testing.assert_allclose(clone()[0], ubq_system.coords) + torch.testing.assert_close( + clone()[0], + torch.tensor(ubq_system.coords, dtype=clone()[0].dtype), + equal_nan=True, + ) clone.dofs[0] -= 1 - torch.testing.assert_allclose(clone()[0], ubq_system.coords) + torch.testing.assert_close( + clone()[0], + torch.tensor(ubq_system.coords, dtype=clone()[0].dtype), + equal_nan=True, + ) # Device can be overridden clone = clone.to(cuda_device) @@ -88,7 +114,11 @@ def test_kinematic_dof_factory(ubq_system): assert clone().device == cuda_device # Coords are returned from forward - torch.testing.assert_allclose(clone().cpu()[0], ubq_system.coords) + torch.testing.assert_close( + clone().cpu()[0], + torch.tensor(ubq_system.coords, dtype=clone().dtype), + equal_nan=True, + ) @pytest.fixture diff --git a/tmol/tests/kinematics/test_gpu_operations.py b/tmol/tests/kinematics/test_gpu_operations.py index 27ca4a98b..1e7768f70 100644 --- a/tmol/tests/kinematics/test_gpu_operations.py +++ b/tmol/tests/kinematics/test_gpu_operations.py @@ -95,7 +95,7 @@ def parallel_refold_hts_cpp(): assert dofs_cuda.device.type == "cuda" assert dofs_cpu.device.type == "cpu" - torch.testing.assert_allclose(dofs_cuda.cpu(), dofs_cpu) + torch.testing.assert_close(dofs_cuda.cpu(), dofs_cpu) @requires_cuda diff --git a/tmol/tests/kinematics/test_script_modules.py b/tmol/tests/kinematics/test_script_modules.py index 6f4d2493d..e805a4101 100644 --- a/tmol/tests/kinematics/test_script_modules.py +++ b/tmol/tests/kinematics/test_script_modules.py @@ -32,7 +32,7 @@ def test_kinematic_torch_op_forward(benchmark, ubq_system, torch_device): def refold_kincoords(): return kop(tdofs.raw) - torch.testing.assert_allclose(refold_kincoords, kincoords) + torch.testing.assert_close(refold_kincoords, kincoords) assert refold_kincoords.device.type == torch_device.type @@ -55,7 +55,7 @@ def test_kinematic_torch_op_backward_benchmark(benchmark, ubq_system, torch_devi def refold_grad(): total.backward(retain_graph=True) - torch.testing.assert_allclose(refold_kincoords, kincoords) + torch.testing.assert_close(refold_kincoords, kincoords) assert refold_kincoords.device.type == torch_device.type diff --git a/tmol/tests/optimization/test_lbfgs_armijo.py b/tmol/tests/optimization/test_lbfgs_armijo.py index 86c0738f3..fed508f13 100755 --- a/tmol/tests/optimization/test_lbfgs_armijo.py +++ b/tmol/tests/optimization/test_lbfgs_armijo.py @@ -10,24 +10,18 @@ def __init__(self, r_m=1.0, epsilon=1.0): self.epsilon = epsilon def __call__(self, coords): - ind = torch.arange(coords.shape[0], requires_grad=False) - - ind_a = ind.view((-1, 1)) - ind_b = ind.view((1, -1)) - deltas = coords.view((-1, 1, 3)) - coords.view((1, -1, 3)) - + N = coords.shape[0] + row, col = torch.tril_indices(N, N, offset=-1) + deltas = coords[row] - coords[col] dist = torch.norm(deltas, 2, -1) - fd = self.r_m / dist + fd2 = fd * fd fd6 = fd2 * fd2 * fd2 fd12 = fd6 * fd6 - lj = self.epsilon * (fd12 - 3 * fd6) - - self.lj = torch.where(ind_a != ind_b, lj, torch.Tensor([0.0])) + self.lj = self.epsilon * (fd12 - 3 * fd6) - self.atom_scores = torch.sum(self.lj.detach(), dim=-1) - self.total_score = torch.sum(self.lj) + self.total_score = 2 * torch.sum(self.lj) return self diff --git a/tmol/tests/pack/rotamer/test_single_residue_kintree.py b/tmol/tests/pack/rotamer/test_single_residue_kintree.py index e339336b3..64bf3059d 100644 --- a/tmol/tests/pack/rotamer/test_single_residue_kintree.py +++ b/tmol/tests/pack/rotamer/test_single_residue_kintree.py @@ -14,18 +14,18 @@ def test_annotate_restypes(default_database): construct_single_residue_kinforest(rt) assert hasattr(rt, "rotamer_kinforest") - assert type(rt.rotamer_kinforest.kinforest_idx) == numpy.ndarray - assert type(rt.rotamer_kinforest.id) == numpy.ndarray - assert type(rt.rotamer_kinforest.doftype) == numpy.ndarray - assert type(rt.rotamer_kinforest.parent) == numpy.ndarray - assert type(rt.rotamer_kinforest.frame_x) == numpy.ndarray - assert type(rt.rotamer_kinforest.frame_y) == numpy.ndarray - assert type(rt.rotamer_kinforest.frame_z) == numpy.ndarray - assert type(rt.rotamer_kinforest.nodes) == numpy.ndarray - assert type(rt.rotamer_kinforest.scans) == numpy.ndarray - assert type(rt.rotamer_kinforest.gens) == numpy.ndarray - assert type(rt.rotamer_kinforest.n_scans_per_gen) == numpy.ndarray - assert type(rt.rotamer_kinforest.dofs_ideal) == numpy.ndarray + assert type(rt.rotamer_kinforest.kinforest_idx) is numpy.ndarray + assert type(rt.rotamer_kinforest.id) is numpy.ndarray + assert type(rt.rotamer_kinforest.doftype) is numpy.ndarray + assert type(rt.rotamer_kinforest.parent) is numpy.ndarray + assert type(rt.rotamer_kinforest.frame_x) is numpy.ndarray + assert type(rt.rotamer_kinforest.frame_y) is numpy.ndarray + assert type(rt.rotamer_kinforest.frame_z) is numpy.ndarray + assert type(rt.rotamer_kinforest.nodes) is numpy.ndarray + assert type(rt.rotamer_kinforest.scans) is numpy.ndarray + assert type(rt.rotamer_kinforest.gens) is numpy.ndarray + assert type(rt.rotamer_kinforest.n_scans_per_gen) is numpy.ndarray + assert type(rt.rotamer_kinforest.dofs_ideal) is numpy.ndarray assert rt.rotamer_kinforest.kinforest_idx.shape == (rt.n_atoms,) assert rt.rotamer_kinforest.id.shape == (rt.n_atoms,) diff --git a/tmol/tests/score/common/geom/test_cuda.py b/tmol/tests/score/common/geom/test_cuda.py index 62ed42396..c6ddfbcd6 100644 --- a/tmol/tests/score/common/geom/test_cuda.py +++ b/tmol/tests/score/common/geom/test_cuda.py @@ -20,5 +20,5 @@ def test_distance(geom): B = torch.ones((10, 3), device="cuda") V = torch.full((10,), math.sqrt(3), device="cuda") - torch.testing.assert_allclose(geom.distance_V(A, B), V) - torch.testing.assert_allclose(geom.distance_V_dV(A, B)[0], V) + torch.testing.assert_close(geom.distance_V(A, B), V) + torch.testing.assert_close(geom.distance_V_dV(A, B)[0], V) diff --git a/tmol/tests/score/common/test_stack_condense.py b/tmol/tests/score/common/test_stack_condense.py index 453786054..97a0154e6 100644 --- a/tmol/tests/score/common/test_stack_condense.py +++ b/tmol/tests/score/common/test_stack_condense.py @@ -30,7 +30,7 @@ def test_condense_torch_inds(torch_device): dtype=torch.int64, device=torch_device, ) - torch.testing.assert_allclose(condensed_inds, expected) + torch.testing.assert_close(condensed_inds, expected) assert condensed_inds.device == torch_device @@ -45,7 +45,7 @@ def test_take_values_w_sentineled_index1(torch_device): dtype=torch.int32, device=torch_device, ) - torch.testing.assert_allclose(index_values, expected) + torch.testing.assert_close(index_values, expected) assert index_values.dtype == torch.int32 assert index_values.device == torch_device @@ -70,7 +70,7 @@ def test_take_values_w_sentineled_index_and_dest(torch_device): device=torch_device, ) - torch.testing.assert_allclose(index_values, expected) + torch.testing.assert_close(index_values, expected) def test_condense_subset(torch_device): @@ -91,7 +91,7 @@ def test_condense_subset(torch_device): ) subset = sc.condense_subset(vals, vals_to_keep) - torch.testing.assert_allclose(subset, expected) + torch.testing.assert_close(subset, expected) def test_condense_numpy_inds_from_doc_string(): @@ -105,7 +105,7 @@ def test_condense_torch_inds_from_doc_string(): input = torch.tensor([[0, 1, 0, 1], [1, 1, 0, 1]], dtype=torch.int32) == 1 expected_output = torch.tensor([[1, 3, -1], [0, 1, 3]], dtype=torch.int64) actual_output = sc.condense_torch_inds(input, torch.device("cpu")) - torch.testing.assert_allclose(actual_output, expected_output) + torch.testing.assert_close(actual_output, expected_output) def test_take_values_w_sentineled_index_from_doc_string(): @@ -117,7 +117,7 @@ def test_take_values_w_sentineled_index_from_doc_string(): [[12, 11, 12, 15, -1], [11, 14, 11, 15, 12]], dtype=torch.int32 ) actual_output = sc.take_values_w_sentineled_index(values, sentineled_index_tensor) - torch.testing.assert_allclose(actual_output, expected_output) + torch.testing.assert_close(actual_output, expected_output) def test_take_values_w_sentineled_index_and_dest_from_doc_string(): @@ -137,7 +137,7 @@ def test_take_values_w_sentineled_index_and_dest_from_doc_string(): actual_output = sc.take_values_w_sentineled_index_and_dest( values, sentineled_index_tensor, sentineled_dest_tensor ) - torch.testing.assert_allclose(actual_output, expected_output) + torch.testing.assert_close(actual_output, expected_output) def test_take_values_w_sentineled_dest_from_doc_string(): @@ -159,7 +159,7 @@ def test_take_values_w_sentineled_dest_from_doc_string(): actual_output = sc.take_values_w_sentineled_dest( values, values_to_take, sentineled_dest_tensor ) - torch.testing.assert_allclose(actual_output, expected_output) + torch.testing.assert_close(actual_output, expected_output) def test_condense_subset_from_doc_string(): @@ -182,7 +182,7 @@ def test_condense_subset_from_doc_string(): dtype=torch.int32, ) actual_output = sc.condense_subset(values, values_to_keep) - torch.testing.assert_allclose(actual_output, expected_output) + torch.testing.assert_close(actual_output, expected_output) def test_take_condensed_3d_subset_from_doc_string(): @@ -210,7 +210,7 @@ def test_take_condensed_3d_subset_from_doc_string(): actual_output = sc.take_condensed_3d_subset( values, condensed_inds_to_keep, condensed_dest_tensor ) - torch.testing.assert_allclose(actual_output, expected_output) + torch.testing.assert_close(actual_output, expected_output) @pytest.mark.parametrize("torch_dtype", [torch.int32, torch.int64]) diff --git a/tmol/tests/score/dunbrack/test_dun_param_resolver.py b/tmol/tests/score/dunbrack/test_dun_param_resolver.py index 8f837583e..dc5a820a5 100644 --- a/tmol/tests/score/dunbrack/test_dun_param_resolver.py +++ b/tmol/tests/score/dunbrack/test_dun_param_resolver.py @@ -638,43 +638,43 @@ def test_stacked_dun_param_resolver_construction(default_database, torch_device) expected = expected_stacked_parameters() - torch.testing.assert_allclose( + torch.testing.assert_close( expected["ndihe_for_res"], dun_params.ndihe_for_res.to("cpu") ) - torch.testing.assert_allclose( + torch.testing.assert_close( expected["dihedral_offset_for_res"], dun_params.dihedral_offset_for_res.to("cpu"), ) - torch.testing.assert_allclose( + torch.testing.assert_close( expected["dihedral_atom_inds"], dun_params.dihedral_atom_inds.to("cpu") ) - torch.testing.assert_allclose( + torch.testing.assert_close( expected["rottable_set_for_res"], dun_params.rottable_set_for_res.to("cpu") ) - torch.testing.assert_allclose( + torch.testing.assert_close( expected["nchi_for_res"], dun_params.nchi_for_res.to("cpu") ) - torch.testing.assert_allclose( + torch.testing.assert_close( expected["nrotameric_chi_for_res"], dun_params.nrotameric_chi_for_res.to("cpu") ) - torch.testing.assert_allclose( + torch.testing.assert_close( expected["rotres2resid"], dun_params.rotres2resid.to("cpu") ) - torch.testing.assert_allclose( + torch.testing.assert_close( expected["prob_table_offset_for_rotresidue"], dun_params.prob_table_offset_for_rotresidue.to("cpu"), ) - torch.testing.assert_allclose( + torch.testing.assert_close( expected["rotmean_table_offset_for_residue"], dun_params.rotmean_table_offset_for_residue.to("cpu"), ) - torch.testing.assert_allclose( + torch.testing.assert_close( expected["rotind2tableind_offset_for_res"], dun_params.rotind2tableind_offset_for_res.to("cpu"), ) - torch.testing.assert_allclose( + torch.testing.assert_close( expected["rotameric_chi_desc"], dun_params.rotameric_chi_desc.to("cpu") ) - torch.testing.assert_allclose( + torch.testing.assert_close( expected["semirotameric_chi_desc"], dun_params.semirotameric_chi_desc.to("cpu") ) diff --git a/tmol/tests/score/elec/test_params.py b/tmol/tests/score/elec/test_params.py index cfbacfa9a..59ba33bfe 100644 --- a/tmol/tests/score/elec/test_params.py +++ b/tmol/tests/score/elec/test_params.py @@ -1,4 +1,5 @@ import torch +import numpy from tmol.system.packed import PackedResidueSystem, PackedResidueSystemStack from tmol.system.score_support import score_method_to_even_weights_dict @@ -67,10 +68,10 @@ def rbpl(system): twoubq_dict.atom_names, ) - torch.testing.assert_allclose( + numpy.testing.assert_allclose( ub40_rbpl, tubq_rbpl[0:1, : ub40_rbpl.shape[1], : ub40_rbpl.shape[1]] ) - torch.testing.assert_allclose(ub60_rbpl, tubq_rbpl[1:2, :]) + numpy.testing.assert_allclose(ub60_rbpl, tubq_rbpl[1:2, :]) def test_jagged_parameter_resolution_part_charges( @@ -103,5 +104,5 @@ def part_char(system): twoubq_dict.res_names, twoubq_dict.atom_names ) - torch.testing.assert_allclose(ub40_pcs, tubq_pcs[0:1, : ub40_pcs.shape[1]]) - torch.testing.assert_allclose(ub60_pcs, tubq_pcs[1:2, :]) + torch.testing.assert_close(ub40_pcs, tubq_pcs[0:1, : ub40_pcs.shape[1]]) + torch.testing.assert_close(ub60_pcs, tubq_pcs[1:2, :]) diff --git a/tmol/tests/score/elec/test_script_modules.py b/tmol/tests/score/elec/test_script_modules.py index cb2b08ce3..1476b8b70 100644 --- a/tmol/tests/score/elec/test_script_modules.py +++ b/tmol/tests/score/elec/test_script_modules.py @@ -141,7 +141,7 @@ def test_elec_intra(default_database, ubq_system, torch_device): val = op(s.tcoords, s.tpcs, s.trbpl) - torch.testing.assert_allclose( + torch.testing.assert_close( val.cpu(), torch.tensor((-135.45822,), dtype=torch.float64), atol=1e-4, @@ -159,7 +159,7 @@ def test_elec_intra3(default_database, ubq_res, torch_device): val = op(s.tcoords, s.tpcs, s.trbpl) - torch.testing.assert_allclose( + torch.testing.assert_close( val.cpu(), torch.tensor((-0.428092,), dtype=torch.float64), atol=1e-4, rtol=1e-2 ) @@ -194,6 +194,6 @@ def test_elec_inter(default_database, ubq_system, torch_device): s.trbpl[:, :part, part:], ) - torch.testing.assert_allclose( + torch.testing.assert_close( val.cpu(), torch.tensor((-53.3774,), dtype=torch.float64), atol=1e-4, rtol=1e-2 ) diff --git a/tmol/tests/score/hbond/test_script_module.py b/tmol/tests/score/hbond/test_script_module.py index 0f7bfaf5b..6c09dc700 100644 --- a/tmol/tests/score/hbond/test_script_module.py +++ b/tmol/tests/score/hbond/test_script_module.py @@ -109,5 +109,5 @@ def test_script_module_scores(default_database, ubq_system, torch_device): cpu_compact_db.global_param_table, ) - torch.testing.assert_allclose(module_score.cpu(), batch_score) + torch.testing.assert_close(module_score.cpu(), batch_score) # Derivative values validated via gradcheck diff --git a/tmol/tests/score/ljlk/potentials/test_compiled_lj_potential.py b/tmol/tests/score/ljlk/potentials/test_compiled_lj_potential.py index 8b0626a35..069fb9db7 100644 --- a/tmol/tests/score/ljlk/potentials/test_compiled_lj_potential.py +++ b/tmol/tests/score/ljlk/potentials/test_compiled_lj_potential.py @@ -119,7 +119,7 @@ def eval_d_lj_d_dist(dist, bonded_path_length=5): # Bonded path length weights ds = torch.linspace(0.0, 8.0, 100) - torch.testing.assert_allclose(eval_lj(ds, 4), eval_lj(ds, 5) * 0.2) + torch.testing.assert_close(eval_lj(ds, 4), eval_lj(ds, 5) * 0.2) assert (eval_lj(ds, 2) == 0.0).all() - torch.testing.assert_allclose(eval_lj(ds), eval_lj_alone(ds)) + torch.testing.assert_close(eval_lj(ds), eval_lj_alone(ds)) diff --git a/tmol/tests/score/ljlk/potentials/test_sphere_overlap.py b/tmol/tests/score/ljlk/potentials/test_sphere_overlap.py index 9bb807912..2211733a2 100644 --- a/tmol/tests/score/ljlk/potentials/test_sphere_overlap.py +++ b/tmol/tests/score/ljlk/potentials/test_sphere_overlap.py @@ -53,7 +53,7 @@ def test_compute_block_spheres(extension): (gold_com.reshape((-1, 3)), gold_sphere_radii.reshape((-1, 1))), dim=1 ).reshape((1, 7, 4)) - torch.testing.assert_allclose(gold_spheres, spheres, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(gold_spheres, spheres, rtol=1e-5, atol=1e-5) @requires_cuda @@ -109,7 +109,7 @@ def double_stack_depth(x): ).reshape((1, 7, 4)) gold_spheres2 = double_stack_depth(gold_spheres) - torch.testing.assert_allclose(gold_spheres2, spheres, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(gold_spheres2, spheres, rtol=1e-5, atol=1e-5) def some_coords(): diff --git a/tmol/tests/score/ljlk/test_script_modules.py b/tmol/tests/score/ljlk/test_script_modules.py index 120ae8785..99c4cdc8d 100644 --- a/tmol/tests/score/ljlk/test_script_modules.py +++ b/tmol/tests/score/ljlk/test_script_modules.py @@ -134,7 +134,7 @@ def test_lj_intra_op(benchmark, default_database, ubq_system, torch_device): def op_val(): return op(s.tcoords, s.ttype, s.tbpl) - torch.testing.assert_allclose( + torch.testing.assert_close( op_val, torch.tensor((expected_dense,)).to(torch_device) ) @@ -145,7 +145,7 @@ def op_full(): return res - torch.testing.assert_allclose( + torch.testing.assert_close( op_full, torch.tensor((expected_dense,)).to(torch_device) ) @@ -181,7 +181,7 @@ def test_lj_intra_op_stacked(benchmark, default_database, torch_device, ubq_syst def op_val(): return op(coords2, atype2, atbpl) - torch.testing.assert_allclose( + torch.testing.assert_close( op_val, torch.tensor(expected_dense).to(torch_device).sum().unsqueeze(0).repeat(2), ) @@ -210,7 +210,7 @@ def test_lj_inter_op(default_database, torch_device, ubq_system): s.tbpl[:, :part, part:], ) - torch.testing.assert_allclose(val, torch.tensor((expected_dense,)).to(torch_device)) + torch.testing.assert_close(val, torch.tensor((expected_dense,)).to(torch_device)) subind = torch.arange(0, s.tcoords.shape[1], 100) @@ -249,7 +249,7 @@ def test_lk_intra_op(benchmark, default_database, ubq_system, torch_device): def op_val(): return op(s.tcoords, s.ttype, s.thvy_at_inds, s.tbpl) - torch.testing.assert_allclose( + torch.testing.assert_close( op_val, torch.tensor((expected_dense,)).to(torch_device) ) @@ -260,7 +260,7 @@ def op_full(): return res - torch.testing.assert_allclose( + torch.testing.assert_close( op_full, torch.tensor((expected_dense,)).to(torch_device) ) @@ -303,7 +303,7 @@ def test_lk_inter_op(default_database, torch_device, ubq_system): s.tbpl[:, :part, part:], ) - torch.testing.assert_allclose(val, torch.tensor((expected_dense,)).to(torch_device)) + torch.testing.assert_close(val, torch.tensor((expected_dense,)).to(torch_device)) subind = torch.arange(0, s.tcoords.shape[1], 100) diff --git a/tmol/tests/score/lk_ball/potentials/test_compiled_lk_ball.py b/tmol/tests/score/lk_ball/potentials/test_compiled_lk_ball.py index f3ff96136..25b82eaf3 100644 --- a/tmol/tests/score/lk_ball/potentials/test_compiled_lk_ball.py +++ b/tmol/tests/score/lk_ball/potentials/test_compiled_lk_ball.py @@ -43,7 +43,7 @@ def test_build_acc_waters(): for torsion, water_ref in zip(torsions, waters_ref): water = BuildAcceptorWater.apply(A, B, B0, dist, angle, torsion) - torch.testing.assert_allclose(water, water_ref) + torch.testing.assert_close(water, water_ref) gradcheck( lambda A, B, B0: BuildAcceptorWater.apply(A, B, B0, dist, angle, torsion), @@ -66,7 +66,7 @@ def test_build_don_water(): waters = BuildDonorWater.apply(D, H, dist) waters_ref = tensor([-7.91642236, 3.81579633, 1.5335272]) - torch.testing.assert_allclose(waters, waters_ref) + torch.testing.assert_close(waters, waters_ref) gradcheck( lambda D, H: BuildDonorWater.apply(D, H, dist), @@ -278,7 +278,7 @@ def test_lk_ball_donor_donor_spotcheck(ljlk_params, atype_params): at_j, ) - torch.testing.assert_allclose( + torch.testing.assert_close( i_by_j + j_by_i, tensor([0.3355, 0.0, 0.2649, 0.7896]), atol=1e-4, rtol=1e-4 ) @@ -333,7 +333,7 @@ def test_lk_ball_sp2_nonpolar_spotcheck(ljlk_params, atype_params): nonpolar_at, ) - torch.testing.assert_allclose(i_by_j, tensor([0.14107985, 0.04765878, 0.0, 0.0])) + torch.testing.assert_close(i_by_j, tensor([0.14107985, 0.04765878, 0.0, 0.0])) def test_lk_ball_sp3_ring_spotcheck(ljlk_params, atype_params): @@ -441,7 +441,7 @@ def test_lk_ball_sp3_ring_spotcheck(ljlk_params, atype_params): atom_types[3], ) - torch.testing.assert_allclose( + torch.testing.assert_close( sp3_by_nonpolar, tensor([0.00385956, 0.0001626, 0.0, 0.0]), atol=1e-4, rtol=1e-4 ) @@ -457,7 +457,7 @@ def test_lk_ball_sp3_ring_spotcheck(ljlk_params, atype_params): atom_types[5], ) - torch.testing.assert_allclose( + torch.testing.assert_close( sp3_by_nonpolar, tensor([0.00369549, 0.0028072, 0.0, 0.0]), atol=1e-4, rtol=1e-4 ) @@ -472,7 +472,7 @@ def test_lk_ball_sp3_ring_spotcheck(ljlk_params, atype_params): ring_at, atom_types[0], ) - torch.testing.assert_allclose( + torch.testing.assert_close( ring_by_nonpolar, tensor([0.01360676, 0.0135272, 0.0, 0.0]), atol=1e-4, diff --git a/tmol/tests/score/lk_ball/test_script_modules.py b/tmol/tests/score/lk_ball/test_script_modules.py index 78f41b0c0..a68f0738f 100644 --- a/tmol/tests/score/lk_ball/test_script_modules.py +++ b/tmol/tests/score/lk_ball/test_script_modules.py @@ -134,7 +134,7 @@ def test_lkball_intra(test_case, torch_device, default_database): indexed_bonds.bond_spans, ) - torch.testing.assert_allclose( + torch.testing.assert_close( val.cpu(), test_case.expected_score, atol=1e-4, rtol=1e-3 ) @@ -211,7 +211,7 @@ def test_lkball_inter(test_case, torch_device, default_database): indexed_bonds.bond_spans, ) - torch.testing.assert_allclose( + torch.testing.assert_close( val.cpu(), test_case.expected_score, atol=1e-4, rtol=1e-3 ) diff --git a/tmol/tests/score/modules/test_cartbonded.py b/tmol/tests/score/modules/test_cartbonded.py index e8d18f0af..def8427da 100644 --- a/tmol/tests/score/modules/test_cartbonded.py +++ b/tmol/tests/score/modules/test_cartbonded.py @@ -86,13 +86,13 @@ def test_cartbonded_for_stacked_system(ubq_system: PackedResidueSystem): tot = stacked_score.intra_total(coords) assert tot.shape == (2,) - torch.testing.assert_allclose(tot[0], tot[1]) + torch.testing.assert_close(tot[0], tot[1]) forward = stacked_score.intra_forward(coords) assert len(forward) == 5 for terms in forward.values(): assert len(terms) == 2 - torch.testing.assert_allclose(terms[0], terms[1]) + torch.testing.assert_close(terms[0], terms[1]) sumtot = torch.sum(tot) sumtot.backward() diff --git a/tmol/tests/score/modules/test_constraint.py b/tmol/tests/score/modules/test_constraint.py index d777f2e98..e6dd5a1c4 100644 --- a/tmol/tests/score/modules/test_constraint.py +++ b/tmol/tests/score/modules/test_constraint.py @@ -150,7 +150,7 @@ def test_cst_for_system(cst_system, cst_csts, torch_device): assert len(tot) == 3 # TODO ask frank what the correct values should be - # torch.testing.assert_allclose(tot.cpu(), -15955.91015625) + # torch.testing.assert_close(tot.cpu(), -15955.91015625) @pytest.mark.benchmark(group="score_components") @@ -191,4 +191,4 @@ def stack_score_constraints(): assert len(tot) == 3 # TODO ask frank what the correct values should be - # torch.testing.assert_allclose(tot.cpu(), -15955.91015625 * nstacks) + # torch.testing.assert_close(tot.cpu(), -15955.91015625 * nstacks) diff --git a/tmol/tests/score/modules/test_dunbrack.py b/tmol/tests/score/modules/test_dunbrack.py index 838763c85..3c7a0e277 100644 --- a/tmol/tests/score/modules/test_dunbrack.py +++ b/tmol/tests/score/modules/test_dunbrack.py @@ -85,13 +85,13 @@ def test_dunbrack_for_stacked_system(ubq_system: PackedResidueSystem): tot = stacked_score.intra_total(coords) assert tot.shape == (2,) - torch.testing.assert_allclose(tot[0], tot[1]) + torch.testing.assert_close(tot[0], tot[1]) forward = stacked_score.intra_forward(coords) assert len(forward) == 3 for terms in forward.values(): assert len(terms) == 2 - torch.testing.assert_allclose(terms[0], terms[1]) + torch.testing.assert_close(terms[0], terms[1]) sumtot = torch.sum(tot) sumtot.backward() diff --git a/tmol/tests/score/modules/test_elec.py b/tmol/tests/score/modules/test_elec.py index cdeb3573a..e929fcd6d 100644 --- a/tmol/tests/score/modules/test_elec.py +++ b/tmol/tests/score/modules/test_elec.py @@ -65,13 +65,13 @@ def test_elec_for_stacked_system(ubq_system: PackedResidueSystem): tot = stacked_score.intra_total(coords) assert tot.shape == (2,) - torch.testing.assert_allclose(tot[0], tot[1]) + torch.testing.assert_close(tot[0], tot[1]) forward = stacked_score.intra_forward(coords) assert len(forward) == 1 for terms in forward.values(): assert len(terms) == 2 - torch.testing.assert_allclose(terms[0], terms[1]) + torch.testing.assert_close(terms[0], terms[1]) sumtot = torch.sum(tot) sumtot.backward() diff --git a/tmol/tests/score/modules/test_hbond.py b/tmol/tests/score/modules/test_hbond.py index 6973adf86..539f934e2 100644 --- a/tmol/tests/score/modules/test_hbond.py +++ b/tmol/tests/score/modules/test_hbond.py @@ -66,13 +66,13 @@ def test_hbond_for_stacked_system(ubq_system: PackedResidueSystem): tot = stacked_score.intra_total(coords) assert tot.shape == (2,) - torch.testing.assert_allclose(tot[0], tot[1]) + torch.testing.assert_close(tot[0], tot[1]) forward = stacked_score.intra_forward(coords) assert len(forward) == 1 for terms in forward.values(): assert len(terms) == 2 - torch.testing.assert_allclose(terms[0], terms[1]) + torch.testing.assert_close(terms[0], terms[1]) sumtot = torch.sum(tot) sumtot.backward() diff --git a/tmol/tests/score/modules/test_ljlk.py b/tmol/tests/score/modules/test_ljlk.py index 71e770e6a..91ccc54f3 100644 --- a/tmol/tests/score/modules/test_ljlk.py +++ b/tmol/tests/score/modules/test_ljlk.py @@ -65,13 +65,13 @@ def test_lj_for_stacked_system(ubq_system: PackedResidueSystem): tot = stacked_score.intra_total(coords) assert tot.shape == (2,) - torch.testing.assert_allclose(tot[0], tot[1]) + torch.testing.assert_close(tot[0], tot[1]) forward = stacked_score.intra_forward(coords) assert len(forward) == 1 for terms in forward.values(): assert len(terms) == 2 - torch.testing.assert_allclose(terms[0], terms[1]) + torch.testing.assert_close(terms[0], terms[1]) sumtot = torch.sum(tot) sumtot.backward() diff --git a/tmol/tests/score/modules/test_lk_ball.py b/tmol/tests/score/modules/test_lk_ball.py index f1393186c..8ffc99756 100644 --- a/tmol/tests/score/modules/test_lk_ball.py +++ b/tmol/tests/score/modules/test_lk_ball.py @@ -73,13 +73,13 @@ def test_lk_ball_for_stacked_system(ubq_system: PackedResidueSystem): tot = stacked_score.intra_total(coords) assert tot.shape == (2,) - torch.testing.assert_allclose(tot[0], tot[1]) + torch.testing.assert_close(tot[0], tot[1]) forward = stacked_score.intra_forward(coords) assert len(forward) == 4 for terms in forward.values(): assert len(terms) == 2 - torch.testing.assert_allclose(terms[0], terms[1]) + torch.testing.assert_close(terms[0], terms[1]) sumtot = torch.sum(tot) sumtot.backward() diff --git a/tmol/tests/score/modules/test_old_bonded_atom.py b/tmol/tests/score/modules/test_old_bonded_atom.py index 50292e7a6..afa1cdb3e 100644 --- a/tmol/tests/score/modules/test_old_bonded_atom.py +++ b/tmol/tests/score/modules/test_old_bonded_atom.py @@ -141,6 +141,6 @@ def test_variable_bonded_path_length(ubq_res): numpy.testing.assert_allclose( inds_both.bond_spans[0, : inds4.bond_spans.shape[1]], inds4.bond_spans[0] ) - torch.testing.assert_allclose( + torch.testing.assert_close( inds_both.bond_spans[1, : inds6.bond_spans.shape[1]], inds6.bond_spans[0] ) diff --git a/tmol/tests/score/modules/test_omega.py b/tmol/tests/score/modules/test_omega.py index b95b61483..c328c56fb 100644 --- a/tmol/tests/score/modules/test_omega.py +++ b/tmol/tests/score/modules/test_omega.py @@ -33,13 +33,13 @@ def test_lj_for_stacked_system(ubq_system: PackedResidueSystem): tot = stacked_score.intra_total(coords) assert tot.shape == (2,) - torch.testing.assert_allclose(tot[0], tot[1]) + torch.testing.assert_close(tot[0], tot[1]) forward = stacked_score.intra_forward(coords) assert len(forward) == 1 for terms in forward.values(): assert len(terms) == 2 - torch.testing.assert_allclose(terms[0], terms[1]) + torch.testing.assert_close(terms[0], terms[1]) sumtot = torch.sum(tot) sumtot.backward() diff --git a/tmol/tests/score/modules/test_rama.py b/tmol/tests/score/modules/test_rama.py index d064fbc01..a4acbe523 100644 --- a/tmol/tests/score/modules/test_rama.py +++ b/tmol/tests/score/modules/test_rama.py @@ -65,13 +65,13 @@ def test_rama_for_stacked_system(ubq_system: PackedResidueSystem): tot = stacked_score.intra_total(coords) assert tot.shape == (2,) - torch.testing.assert_allclose(tot[0], tot[1]) + torch.testing.assert_close(tot[0], tot[1]) forward = stacked_score.intra_forward(coords) assert len(forward) == 1 for terms in forward.values(): assert len(terms) == 2 - torch.testing.assert_allclose(terms[0], terms[1]) + torch.testing.assert_close(terms[0], terms[1]) sumtot = torch.sum(tot) sumtot.backward() diff --git a/tmol/tests/support/test_database_converters.py b/tmol/tests/support/test_database_converters.py index 5480ecef3..f1bfb6ecf 100644 --- a/tmol/tests/support/test_database_converters.py +++ b/tmol/tests/support/test_database_converters.py @@ -94,20 +94,20 @@ def test_dunbrack_table_read(rosetta_database, default_database): assert len(default.semi_rotameric_libraries) == len(fresh.semi_rotameric_libraries) def compare_rotameric_data(default_rotdat, fresh_rotdat): - torch.testing.assert_allclose(default_rotdat.rotamers, fresh_rotdat.rotamers) - torch.testing.assert_allclose( + torch.testing.assert_close(default_rotdat.rotamers, fresh_rotdat.rotamers) + torch.testing.assert_close( default_rotdat.rotamer_probabilities, fresh_rotdat.rotamer_probabilities ) - torch.testing.assert_allclose( + torch.testing.assert_close( default_rotdat.rotamer_means, fresh_rotdat.rotamer_means ) - torch.testing.assert_allclose( + torch.testing.assert_close( default_rotdat.backbone_dihedral_start, fresh_rotdat.backbone_dihedral_start ) - torch.testing.assert_allclose( + torch.testing.assert_close( default_rotdat.backbone_dihedral_step, fresh_rotdat.backbone_dihedral_step ) - torch.testing.assert_allclose( + torch.testing.assert_close( default_rotdat.rotamer_alias, fresh_rotdat.rotamer_alias ) @@ -123,23 +123,23 @@ def compare_rotameric_data(default_rotdat, fresh_rotdat): def compare_semirotameric_data(default_srdat, fresh_srdat): compare_rotameric_data(default_srdat.rotameric_data, fresh_srdat.rotameric_data) - torch.testing.assert_allclose( + torch.testing.assert_close( default_srdat.non_rot_chi_start, fresh_srdat.non_rot_chi_start ) - torch.testing.assert_allclose( + torch.testing.assert_close( default_srdat.non_rot_chi_step, fresh_srdat.non_rot_chi_step ) - torch.testing.assert_allclose( + torch.testing.assert_close( default_srdat.non_rot_chi_period, fresh_srdat.non_rot_chi_period ) - torch.testing.assert_allclose( + torch.testing.assert_close( default_srdat.rotameric_chi_rotamers, fresh_srdat.rotameric_chi_rotamers ) - torch.testing.assert_allclose( + torch.testing.assert_close( default_srdat.nonrotameric_chi_probabilities, fresh_srdat.nonrotameric_chi_probabilities, ) - torch.testing.assert_allclose( + torch.testing.assert_close( default_srdat.rotamer_boundaries, fresh_srdat.rotamer_boundaries ) diff --git a/tmol/tests/test_cuda.py b/tmol/tests/test_cuda.py index b9e850d4e..554d0b32e 100644 --- a/tmol/tests/test_cuda.py +++ b/tmol/tests/test_cuda.py @@ -20,7 +20,7 @@ def test_torch_cuda_smoke(): c = a.cuda() @ b.cuda() - torch.testing.assert_allclose(a @ b, c.cpu()) + torch.testing.assert_close(a @ b, c.cpu()) @pytest.mark.parametrize("dtype", [torch.float, torch.double], ids=("single", "double")) diff --git a/tmol/tests/utility/tensor/test_tensor_accessor.py b/tmol/tests/utility/tensor/test_tensor_accessor.py index 6a6ef72fa..154d267f8 100644 --- a/tmol/tests/utility/tensor/test_tensor_accessor.py +++ b/tmol/tests/utility/tensor/test_tensor_accessor.py @@ -29,7 +29,7 @@ def test_tensor_vector_accessors(accessor_funcs): results = {t: f(tvec) for t, f in accessor_funcs.items()} for _rn, r in results.items(): - torch.testing.assert_allclose(r, expected) + torch.testing.assert_close(r, expected) @pytest.fixture @@ -46,7 +46,7 @@ def test_tensor_matrix_accessors(matrix_accessor_funcs): results = {t: f(tvec) for t, f in matrix_accessor_funcs.items()} for _rn, r in results.items(): - torch.testing.assert_allclose(r, expected) + torch.testing.assert_close(r, expected) @requires_cuda @@ -71,25 +71,25 @@ def test_tensor_accessor_device_conversion(accessor_funcs): with pytest.raises(errors[n]): f(tvec) else: - torch.testing.assert_allclose(f(tvec), expected) + torch.testing.assert_close(f(tvec), expected) def test_tensor_pack_eigen_matrix(tensor_accessor): eshape = (2, 5, 3, 3) res = tensor_accessor.tensor_pack_construct_eigen_matrix() - torch.testing.assert_allclose(res[1], torch.ones(eshape)) - torch.testing.assert_allclose(res[2], torch.zeros(eshape)) - torch.testing.assert_allclose(res[3], torch.full(eshape, math.nan)) + torch.testing.assert_close(res[1], torch.ones(eshape)) + torch.testing.assert_close(res[2], torch.zeros(eshape)) + torch.testing.assert_close(res[3], torch.full(eshape, math.nan), equal_nan=True) def test_tensor_pack_constructors(tensor_accessor): eshape = (2, 5, 3) res = tensor_accessor.tensor_pack_construct() - torch.testing.assert_allclose(res[1], torch.ones(eshape)) - torch.testing.assert_allclose(res[2], torch.zeros(eshape)) - torch.testing.assert_allclose(res[3], torch.full(eshape, math.nan)) + torch.testing.assert_close(res[1], torch.ones(eshape)) + torch.testing.assert_close(res[2], torch.zeros(eshape)) + torch.testing.assert_close(res[3], torch.full(eshape, math.nan), equal_nan=True) t = torch.empty((1, 4)) eshape = (1, 4, 3) @@ -97,9 +97,9 @@ def test_tensor_pack_constructors(tensor_accessor): # ATen *_like constructors res = tensor_accessor.tensor_pack_construct_like_aten(t) - torch.testing.assert_allclose(res[1], torch.ones(eshape)) - torch.testing.assert_allclose(res[2], torch.zeros(eshape)) - torch.testing.assert_allclose(res[3], torch.full(eshape, math.nan)) + torch.testing.assert_close(res[1], torch.ones(eshape)) + torch.testing.assert_close(res[2], torch.zeros(eshape)) + torch.testing.assert_close(res[3], torch.full(eshape, math.nan), equal_nan=True) with pytest.raises(RuntimeError): tensor_accessor.tensor_pack_construct_like_aten(torch.empty(10)) @@ -107,9 +107,9 @@ def test_tensor_pack_constructors(tensor_accessor): # TView *_like constructors res = tensor_accessor.tensor_pack_construct_like_tview(t) - torch.testing.assert_allclose(res[1], torch.ones(eshape)) - torch.testing.assert_allclose(res[2], torch.zeros(eshape)) - torch.testing.assert_allclose(res[3], torch.full(eshape, math.nan)) + torch.testing.assert_close(res[1], torch.ones(eshape)) + torch.testing.assert_close(res[2], torch.zeros(eshape)) + torch.testing.assert_close(res[3], torch.full(eshape, math.nan), equal_nan=True) with pytest.raises(TypeError): tensor_accessor.tensor_pack_construct_like_tview(torch.empty(10)) @@ -117,9 +117,9 @@ def test_tensor_pack_constructors(tensor_accessor): # TPack *_like constructors res = tensor_accessor.tensor_pack_construct_like_tpack(t) - torch.testing.assert_allclose(res[1], torch.ones(eshape)) - torch.testing.assert_allclose(res[2], torch.zeros(eshape)) - torch.testing.assert_allclose(res[3], torch.full(eshape, math.nan)) + torch.testing.assert_close(res[1], torch.ones(eshape)) + torch.testing.assert_close(res[2], torch.zeros(eshape)) + torch.testing.assert_close(res[3], torch.full(eshape, math.nan), equal_nan=True) with pytest.raises(TypeError): tensor_accessor.tensor_pack_construct_like_tpack(torch.empty(10)) @@ -128,4 +128,4 @@ def test_tensor_pack_constructors(tensor_accessor): def test_tview_slice(tensor_accessor): sliced = tensor_accessor.tensor_view_take_slice_one() gold = torch.tensor([5, 15, 25, 35], dtype=torch.int32) - torch.testing.assert_allclose(sliced, gold) + torch.testing.assert_close(sliced, gold) diff --git a/tmol/tests/utility/tensor/test_tensor_collection.py b/tmol/tests/utility/tensor/test_tensor_collection.py index af575025d..ebaa7c5fd 100644 --- a/tmol/tests/utility/tensor/test_tensor_collection.py +++ b/tmol/tests/utility/tensor/test_tensor_collection.py @@ -19,4 +19,4 @@ def test_tensor_collection(tensor_collection): ] tsum = tensor_collection.sum_tensor_collection(tcoll) expected = tcoll[0] + tcoll[1] - torch.testing.assert_allclose(expected, tsum) + torch.testing.assert_close(expected, tsum) diff --git a/tmol/tests/utility/torchscript/test_custom_op.py b/tmol/tests/utility/torchscript/test_custom_op.py index 0c6f0b33f..9367a2610 100644 --- a/tmol/tests/utility/torchscript/test_custom_op.py +++ b/tmol/tests/utility/torchscript/test_custom_op.py @@ -4,7 +4,7 @@ def test_load(): # Initial fetch of op fails with RuntimeError, op not registered - with pytest.raises(RuntimeError): + with pytest.raises(AttributeError): torch.ops.tmol.cpow from .custom_op import cpow @@ -17,7 +17,7 @@ def check_form(pow3_f): result = pow3_f(i) expected = i.pow(3.0) - torch.testing.assert_allclose(result, expected) + torch.testing.assert_close(result, expected) assert not result.requires_grad assert not expected.requires_grad @@ -27,7 +27,7 @@ def check_form(pow3_f): result = pow3_f(i) expected = i.pow(3.0) - torch.testing.assert_allclose(result, expected) + torch.testing.assert_close(result, expected) assert result.requires_grad assert expected.requires_grad @@ -41,7 +41,7 @@ def check_form(pow3_f): pow3_f(i).sum().backward() result_grad = i.grad - torch.testing.assert_allclose(result_grad, expected_grad) + torch.testing.assert_close(result_grad, expected_grad) def pow3(t): return cpow(t, 3.0) diff --git a/tmol/types/subscriptable.py b/tmol/types/subscriptable.py new file mode 100644 index 000000000..ee925ee13 --- /dev/null +++ b/tmol/types/subscriptable.py @@ -0,0 +1,70 @@ +class _SubscribedType(type): + """ + This class is a placeholder to let the IDE know the attributes of the + returned type after a __getitem__. + """ + + __origin__ = None + __args__ = None + + +# adapted from typish (https://github.com/ramonhagenaars/typish) +class SubscriptableType(type): + """ + This metaclass will allow a type to become subscriptable. + + >>> class SomeType(metaclass=SubscriptableType): + ... pass + >>> SomeTypeSub = SomeType['some args'] + >>> SomeTypeSub.__args__ + 'some args' + >>> SomeTypeSub.__origin__.__name__ + 'SomeType' + """ + + def __init_subclass__(mcs, **kwargs): + mcs._hash = None + mcs.__args__ = None + mcs.__origin__ = None + + def __getitem__(self, item) -> _SubscribedType: + body = { + **self.__dict__, + "__args__": item, + "__origin__": self, + } + bases = self, *self.__bases__ + result = type(self.__name__, bases, body) + if hasattr(result, "_after_subscription"): + # TODO check if _after_subscription is static + result._after_subscription(item) + return result + + def __eq__(self, other): + self_module = getattr(self, "__module__", None) + self_qualname = getattr(self, "__qualname__", None) + self_origin = getattr(self, "__origin__", None) + self_args = getattr(self, "__args__", None) + + other_module = getattr(other, "__module__", None) + other_qualname = getattr(other, "__qualname__", None) + other_args = getattr(other, "__args__", None) + other_origin = getattr(other, "__origin__", None) + + return ( + self_module == other_module + and self_qualname == other_qualname + and self_args == other_args + and self_origin == other_origin + ) + + def __hash__(self): + if not getattr(self, "_hash", None): + self_module = getattr(self, "__module__", None) + self_qualname = getattr(self, "__qualname__", None) + self_origin = getattr(self, "__origin__", None) + self_args = getattr(self, "__args__", None) + self._hash = hash( + "{}{}{}{}".format(self_module, self_qualname, self_origin, self_args) + ) + return self._hash diff --git a/tmol/types/tensor.py b/tmol/types/tensor.py index bb0ce8085..48aee408e 100644 --- a/tmol/types/tensor.py +++ b/tmol/types/tensor.py @@ -11,7 +11,7 @@ from .converters import register_converter from .validators import register_validator -from typish import SubscriptableType +from .subscriptable import SubscriptableType _NOTHING = object() diff --git a/tmol/utility/cpp_extension.py b/tmol/utility/cpp_extension.py index 478340761..55386d723 100644 --- a/tmol/utility/cpp_extension.py +++ b/tmol/utility/cpp_extension.py @@ -21,7 +21,7 @@ _default_include_paths = list(tmol_include_paths() + extern_include_paths()) -_required_flags = ["--std=c++14", "-DWITH_NVTX", "-w"] +_required_flags = ["--std=c++17", "-DWITH_NVTX", "-w"] if os.environ.get("DEBUG"): _default_flags = ["-O3", "-DDEBUG"] @@ -42,7 +42,7 @@ def get_torch_version(): torch_major, torch_minor = get_torch_version() _required_cuda_flags = [ - "-std=c++14", + "-std=c++17", "--expt-extended-lambda", "-DWITH_NVTX", "-w", diff --git a/tmol/utility/reactive.py b/tmol/utility/reactive.py index 6b871bf7d..d14506b6a 100644 --- a/tmol/utility/reactive.py +++ b/tmol/utility/reactive.py @@ -397,6 +397,7 @@ def grudge( import inspect from collections import defaultdict from typing import Callable, Any, Optional, Tuple, Union +import sys import types @@ -406,7 +407,91 @@ def grudge( # import tmol.utility.nvtx -def _code( +def _code_py311_312( + argcount, + posonlyargcount, + kwonlyargcount, + nlocals, + stacksize, + flags, + codestring, + constants, + names, + varnames, + filename, + name, + qualname, + firstlineno, + linetable, + exceptiontable, + freevars, + cellvars, + **kwargs, +): + """Construct type.CodeType, ignoring unneeded kwargs.""" + return types.CodeType( + argcount, + posonlyargcount, + kwonlyargcount, + nlocals, + stacksize, + flags, + codestring, + constants, + names, + varnames, + filename, + name, + qualname, + firstlineno, + linetable, + exceptiontable, + freevars, + cellvars, + ) + + +def _code_py310( + argcount, + posonlyargcount, + kwonlyargcount, + nlocals, + stacksize, + flags, + codestring, + constants, + names, + varnames, + filename, + name, + firstlineno, + linetable, + freevars, + cellvars, + **kwargs, +): + """Construct type.CodeType, ignoring unneeded kwargs.""" + return types.CodeType( + argcount, + posonlyargcount, + kwonlyargcount, + nlocals, + stacksize, + flags, + codestring, + constants, + names, + varnames, + filename, + name, + firstlineno, + linetable, + freevars, + cellvars, + ) + + +def _code_py39( argcount, posonlyargcount, kwonlyargcount, @@ -446,6 +531,19 @@ def _code( ) +def _code(**kwargs): + if sys.version_info[:2] == (3, 11) or sys.version_info[:2] == (3, 12): + return _code_py311_312(**kwargs) + elif sys.version_info[:2] == (3, 10): + return _code_py310(**kwargs) + elif sys.version_info[:2] == (3, 9): + return _code_py39(**kwargs) + else: + raise NotImplementedError( + f"_code not implemented for python {sys.version_info}" + ) + + def _code_attrs(c): """Unpack types.CodeType attrs ('co_') into __init__ params ('').""" cattrs = {n[3:]: getattr(c, n) for n in dir(c) if n.startswith("co_")}