Add batching support to sdk

azure-quantum/azure/quantum/job/base_job.py

@@ -149,6 +149,95 @@ def from_input_data(
             **kwargs
         )
 
+    @classmethod
+    def from_input_data_container(
+        cls,
+        workspace: "Workspace",
+        name: str,
+        target: str,
+        input_data: Dict[str, bytes],
+        content_type: ContentType = ContentType.json,
+        blob_name: str = "inputData",
+        encoding: str = "",
+        job_id: str = None,
+        container_name: str = None,
+        provider_id: str = None,
+        input_data_format: str = None,
+        output_data_format: str = None,
+        input_params: Dict[str, Any] = None,
+        session_id: Optional[str] = None,
+        **kwargs
+    ) -> "BaseJob":
+        """Create a new Azure Quantum job based on a list of input_data.
+
+        :param workspace: Azure Quantum workspace to submit the input_data to
+        :type workspace: Workspace
+        :param name: Name of the job
+        :type name: str
+        :param target: Azure Quantum target
+        :type target: str
+        :param input_data: Raw input data to submit
+        :type input_data: Dict
+        :param blob_name: Dict of Input data where the key is the blob name
+        :type blob_name: str
+        :param content_type: Content type, e.g. "application/json"
+        :type content_type: ContentType
+        :param encoding: input_data encoding, e.g. "gzip", defaults to empty string
+        :type encoding: str
+        :param job_id: Job ID, defaults to None
+        :type job_id: str
+        :param container_name: Container name, defaults to None
+        :type container_name: str
+        :param provider_id: Provider ID, defaults to None
+        :type provider_id: str
+        :param input_data_format: Input data format, defaults to None
+        :type input_data_format: str
+        :param output_data_format: Output data format, defaults to None
+        :type output_data_format: str
+        :param input_params: Input parameters, defaults to None
+        :type input_params: Dict[str, Any]
+        :param input_params: Input params for job
+        :type input_params: Dict[str, Any]
+        :return: Azure Quantum Job
+        :rtype: Job
+        """
+        # Generate job ID if not specified
+        if job_id is None:
+            job_id = cls.create_job_id()
+
+        # Create container if it does not yet exist
+        container_uri = workspace.get_container_uri(
+            job_id=job_id,
+            container_name=container_name
+        )
+        logger.debug(f"Container URI: {container_uri}")
+
+        # Upload data to container
+        for blob_name, input_data_item in input_data.items():
+            input_data_uri = cls.upload_input_data(
+                container_uri=container_uri,
+                input_data=input_data_item,
+                content_type=content_type,
+                blob_name=blob_name,
+                encoding=encoding,
+            )
+
+        # Create and submit job
+        return cls.from_storage_uri(
+            workspace=workspace,
+            job_id=job_id,
+            target=target,
+            input_data_uri=input_data_uri,
+            container_uri=container_uri,
+            name=name,
+            input_data_format=input_data_format,
+            output_data_format=output_data_format,
+            provider_id=provider_id,
+            input_params=input_params,
+            session_id=session_id,
+            **kwargs
+        )
+
     @classmethod
     def from_storage_uri(
         cls,

azure-quantum/azure/quantum/target/microsoft/elements/dft/target.py

@@ -5,8 +5,10 @@
 from azure.quantum.target.target import Target
 from azure.quantum.workspace import Workspace
 from azure.quantum.target.params import InputParams
-from typing import Any, Dict, Type, Union
+from typing import Any, Dict, Type, Union, List
 from .job import MicrosoftElementsDftJob
+from pathlib import Path
+import copy
 
 
 class MicrosoftElementsDft(Target):
@@ -73,14 +75,121 @@ def submit(self,
         if shots is not None:
             warnings.warn("The 'shots' parameter is ignored in Microsoft Elements Dft job.")
 
-        return super().submit(
-            input_data=input_data,
-            name=name, 
-            shots=shots, 
-            input_params=input_params,
-            **kwargs
-        )
+        if isinstance(input_data, list):
+
+            qcschema_data = self._assemble_qcshema_from_files(input_data, input_params)
+
+            qcschema_blobs = {}
+            for i in range(len(qcschema_data)):
+                qcschema_blobs[f"input_data_{i}.json"] = self._encode_input_data(qcschema_data[i])
+
+            return self._get_job_class().from_input_data_container(
+                workspace=self.workspace,
+                name=name,
+                target=self.name,
+                input_data=qcschema_blobs,
+                input_params={ 'number_of_molecules': len(qcschema_data), "input_files": list(qcschema_blobs.keys()), **input_params },
+                content_type=kwargs.pop('content_type', self.content_type),
+                encoding=kwargs.pop('encoding', self.encoding),
+                provider_id=self.provider_id,
+                input_data_format=kwargs.pop('input_data_format', 'microsoft.qc-schema.v1'),
+                output_data_format=kwargs.pop('output_data_format', self.output_data_format),
+                session_id=self.get_latest_session_id(),
+                **kwargs
+            )
+        else:
+            return super().submit(
+                input_data=input_data,
+                name=name, 
+                shots=shots, 
+                input_params=input_params,
+                **kwargs
+            )
+
+
+
+    @classmethod
+    def _assemble_qcshema_from_files(self, input_data: List[str], input_params: Dict) -> str:
+        """
+        Convert a list of files to a list of qcshema objects serialized in json.
+        """
+
+        qcshema_objects = []
+        for file in input_data:
+            file_path = Path(file)
+            if not file_path.exists():
+                raise FileNotFoundError(f"File {file} does not exist.")
+
+            file_data = file_path.read_text()
+            mol = self._xyz_to_qcschema_mol(file_data)
+            new_qcschema = self._new_qcshema( input_params, mol )
+            qcshema_objects.append(new_qcschema)
+
+        return qcshema_objects
+
+    @classmethod
+    def _new_qcshema( self, input_params: Dict[str,Any], mol: Dict[str,Any],  ) -> Dict[str, Any]:
+        """
+        Create a new default qcshema object.
+        """
+
+        if input_params.get("driver") == "go":
+            copy_input_params = copy.deepcopy(input_params)
+            copy_input_params["driver"] = "gradient"
+            new_object = {
+                "schema_name": "qcschema_optimization_input",
+                "schema_version": 1,
+                "initial_molecule": mol,
+            }
+            if copy_input_params.get("keywords") and copy_input_params["keywords"].get("geometryOptimization"):
+                new_object["keywords"] = copy_input_params["keywords"].pop("geometryOptimization")
+            new_object["input_specification"] = copy_input_params
+            return new_object
+        elif input_params.get("driver") == "bomd":
+            copy_input_params = copy.deepcopy(input_params)
+            copy_input_params["driver"] = "gradient"
+            new_object = {
+                "schema_name": "madft_molecular_dynamics_input",
+                "schema_version": 1,
+                "initial_molecule": mol,
+            }
+            if copy_input_params.get("keywords") and copy_input_params["keywords"].get("molecularDynamics"):
+                new_object["keywords"] = copy_input_params["keywords"].pop("molecularDynamics")
+            new_object["input_specification"] = copy_input_params
+            return new_object
+        else:
+            new_object = copy.deepcopy(input_params)
+            new_object.update({
+                "schema_name": "qcschema_input",
+                "schema_version": 1,
+                "molecule": mol,
+            })
+            return new_object
+
+
+    @classmethod
+    def _xyz_to_qcschema_mol(self, file_data: str ) -> Dict[str, Any]:
+        """
+        Convert xyz format to qcschema molecule.
+        """
+
+        lines = file_data.split("\n")
+        assert len(lines) >= 3, "Invalid xyz format."
+        n_atoms = int(lines.pop(0))
+        comment = lines.pop(0)
+        mol = {
+            "geometry": [],
+            "symbols": [],
+        }
+        for line in lines:
+            elements = line.split()
+            if len(elements) != 4:
+                raise ValueError("Invalid xyz format.")
+            symbol, x, y, z = elements
+            mol["symbols"].append(symbol)
+            mol["geometry"].append([float(x), float(y), float(z)])
 
+        return mol
 
     @classmethod
     def _get_job_class(cls) -> Type[Job]:

azure-quantum/tests/unit/test_microsoft_elements_dft.py

@@ -4,6 +4,10 @@
 from common import QuantumTestBase, DEFAULT_TIMEOUT_SECS
 from azure.quantum import JobStatus
 from azure.quantum.job import JobFailedWithResultsError
+from azure.quantum.target.microsoft.elements.dft import MicrosoftElementsDft
+from pytest_regressions import data_regression
+from pathlib import Path
+
 
 @pytest.mark.live_test
 class TestMicrosoftElementsDftJob(QuantumTestBase):
@@ -83,3 +87,112 @@ def _run_job(self, input_params) -> Job:
             job.refresh()
 
             return job
+
+test_file = Path(__file__).parent / "molecule.xyz"
+
+@pytest.mark.parametrize(
+        'input_params', [
+        {
+            "driver": "energy",
+            "model": { "method": "m06-2x", "basis": "def2-svp" },
+        },
+        {
+            "driver": "gradient",
+            "model": { "method": "m06-2x", "basis": "def2-svp" },
+        },
+        {
+            "driver": "hessian",
+            "model": { "method": "m06-2x", "basis": "def2-svp" },
+        },
+        {
+            "driver": "energy",
+            "model": { "method": "m06-2x", "basis": "def2-svp" },
+            "keywords": {
+                "scf": { "method": "rks", "maxSteps": 100, "convergeThreshold": 1e-8, "requireWaveFunction": True},
+                "xcFunctional": { "gridLevel": 3 }
+            },
+        },
+        {
+            "driver": "gradient",
+            "model": { "method": "m06-2x", "basis": "def2-svp" },
+            "keywords": {
+                "scf": { "method": "rks", "maxSteps": 100, "convergeThreshold": 1e-8, "requireWaveFunction": True},
+                "xcFunctional": { "gridLevel": 3 }
+            },
+        },
+        {
+            "driver": "hessian",
+            "model": { "method": "m06-2x", "basis": "def2-svp" },
+            "keywords": {
+                "scf": { "method": "rks", "maxSteps": 100, "convergeThreshold": 1e-8, "requireWaveFunction": True},
+                "xcFunctional": { "gridLevel": 3 }
+            },
+        },
+    ]
+)
+@pytest.mark.parametrize(
+    'input_data', [
+        [ test_file ],
+        [ test_file, test_file ],
+    ]
+)
+def test_assemble_true_qcschema_from_files_success(data_regression, input_params, input_data):
+    target = MicrosoftElementsDft
+    qcschema_data = target._assemble_qcshema_from_files(input_data, input_params)
+    data_regression.check(qcschema_data)
+
+@pytest.mark.parametrize(
+    'input_params', [
+        {
+            "driver": "go",
+            "model": { "method": "m06-2x", "basis": "def2-svp" },
+        },
+        {
+            "driver": "go",
+            "model": { "method": "m06-2x", "basis": "def2-svp" },
+            "keywords": {
+                "scf": { "method": "rks", "maxSteps": 100, "convergeThreshold": 1e-8, "requireWaveFunction": True},
+                "xcFunctional": { "gridLevel": 3 },
+                "geometryOptimization": {"convergence_grms": 0.001667, "convergence_gmax": 0.0025, "convergence_drms": 0.006667, "convergence_dmax":0.01 }
+            },
+        },
+    ]
+)
+@pytest.mark.parametrize(
+    'input_data', [
+        [ test_file ],
+        [ test_file, test_file ],
+    ]
+)
+def test_assemble_go_qcschema_from_files_success(data_regression, input_params, input_data):
+    target = MicrosoftElementsDft
+    qcschema_data = target._assemble_qcshema_from_files(input_data, input_params)
+    data_regression.check(qcschema_data)
+
+@pytest.mark.parametrize(
+    'input_params', [
+        {
+            "driver": "bomd",
+            "model": { "method": "m06-2x", "basis": "def2-svp" },
+        },
+        {
+            "driver": "bomd",
+            "model": { "method": "m06-2x", "basis": "def2-svp" },
+            "keywords": {
+                "scf": { "method": "rks", "maxSteps": 100, "convergeThreshold": 1e-8, "requireWaveFunction": True},
+                "xcFunctional": { "gridLevel": 3 },
+                "molecularDynamics":{"steps": 5, "temperature": 298, "timeStep": 1, "thermostat": {"type": "berendsen", "timeSmoothingFactor": 0.05 } }
+            },
+        },
+    ]
+)
+@pytest.mark.parametrize(
+    'input_data', [
+        [ test_file ],
+        [ test_file, test_file ],
+    ]
+)
+def test_assemble_bomd_qcschema_from_files_success(data_regression, input_params, input_data):
+    target = MicrosoftElementsDft
+    qcschema_data = target._assemble_qcshema_from_files(input_data, input_params)
+    data_regression.check(qcschema_data)