From cd5c78ba1bddf0523f080598166d0d25c0f1f0b6 Mon Sep 17 00:00:00 2001
From: Michele Ceriotti <michele.ceriotti@gmail.com>
Date: Mon, 2 Feb 2026 16:56:06 +0100
Subject: [PATCH 1/3] Retry downloads from zenodo when failing because of error
 423

---
 pyproject.toml                |  1 +
 src/shiftml/ase/calculator.py | 40 +++++++++++++++++++++++------------
 2 files changed, 28 insertions(+), 13 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index c2f6285..2232f50 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -12,6 +12,7 @@ dependencies = [
     "metatensor-torch >=0.7.6,<0.9",
     "metatomic-torch >=0.1.2,<0.2",
     "vesin",
+    "requests",
 ]
 
 readme = "README.md"
diff --git a/src/shiftml/ase/calculator.py b/src/shiftml/ase/calculator.py
index 7ef6e1e..5f37fda 100644
--- a/src/shiftml/ase/calculator.py
+++ b/src/shiftml/ase/calculator.py
@@ -1,6 +1,9 @@
 import logging
 import os
-import urllib.request
+
+import requests
+from requests.adapters import HTTPAdapter
+from urllib3.util.retry import Retry
 
 import numpy as np
 from metatomic.torch import ModelOutput
@@ -57,6 +60,25 @@ def is_fitted_on(atoms, fitted_species):
         )
 
 
+def download_with_retry(url, destination):
+    """Helper function to download data with retries on errors."""
+
+    # Retry strategy: wait 1s, 2s, 4s, 8s, 16s on 429/5xx errors
+    retry_strategy = Retry(
+        total=5, backoff_factor=1, status_forcelist=[429, 500, 502, 503, 504]
+    )
+    session = requests.Session()
+    session.mount("https://", HTTPAdapter(max_retries=retry_strategy))
+
+    # Fetch with automatic retry and error raising
+    response = session.get(url, stream=True)
+    response.raise_for_status()
+
+    with open(destination, "wb") as file:
+        for chunk in response.iter_content(chunk_size=8192):
+            file.write(chunk)
+
+
 def ShiftML(model_version, force_download=False, device=None):
     """
     Initialize the ShiftML calculator
@@ -247,23 +269,15 @@ def __init__(self, model_version, force_download=False, device=None):
                     download = True
 
             if download:
-                urllib.request.urlretrieve(url, model_file)
+                download_with_retry(url, model_file)
                 logging.info(
                     "Downloaded {} and saved to {}".format(model_version, cachedir)
                 )
 
-        except urllib.error.URLError as e:
-            logging.error(
-                "Failed to download {} from {}. URL Error: {}".format(
-                    model_version, url, e.reason
-                )
-            )
-            raise e
-        except urllib.error.HTTPError as e:
+        except requests.exceptions.RequestException as e:
             logging.error(
-                "Failed to download {} from {}.\
-                  HTTP Error: {} - {}".format(
-                    model_version, url, e.code, e.reason
+                "Failed to download {} from {}. Error: {}".format(
+                    model_version, url, e
                 )
             )
             raise e

From e5124d7c36c903d0fae8d5188d42829d8e5ddb67 Mon Sep 17 00:00:00 2001
From: Michele Ceriotti <michele.ceriotti@gmail.com>
Date: Mon, 2 Feb 2026 20:38:37 +0100
Subject: [PATCH 2/3] Lint

---
 src/shiftml/ase/calculator.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/src/shiftml/ase/calculator.py b/src/shiftml/ase/calculator.py
index 5f37fda..0aa65ad 100644
--- a/src/shiftml/ase/calculator.py
+++ b/src/shiftml/ase/calculator.py
@@ -276,9 +276,7 @@ def __init__(self, model_version, force_download=False, device=None):
 
         except requests.exceptions.RequestException as e:
             logging.error(
-                "Failed to download {} from {}. Error: {}".format(
-                    model_version, url, e
-                )
+                "Failed to download {} from {}. Error: {}".format(model_version, url, e)
             )
             raise e
         except Exception as e:

From d83e440d6397f63083050604569ab0405b81b15a Mon Sep 17 00:00:00 2001
From: Michele Ceriotti <michele.ceriotti@gmail.com>
Date: Mon, 2 Feb 2026 22:54:01 +0100
Subject: [PATCH 3/3] more linting

---
 src/shiftml/ase/calculator.py | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/src/shiftml/ase/calculator.py b/src/shiftml/ase/calculator.py
index 0aa65ad..9523416 100644
--- a/src/shiftml/ase/calculator.py
+++ b/src/shiftml/ase/calculator.py
@@ -1,14 +1,13 @@
 import logging
 import os
 
-import requests
-from requests.adapters import HTTPAdapter
-from urllib3.util.retry import Retry
-
 import numpy as np
+import requests
 from metatomic.torch import ModelOutput
 from metatomic.torch.ase_calculator import MetatomicCalculator
 from platformdirs import user_cache_path
+from requests.adapters import HTTPAdapter
+from urllib3.util.retry import Retry
 
 from shiftml.utils.tensorial import T_sym_np_inv, symmetrize