mindspore2022/mindspore/ops/operations/quantum_ops.py

# Copyright 2021 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""Operators for quantum computing."""

from ..primitive import PrimitiveWithInfer, prim_attr_register
from ..._checkparam import Validator as validator
from ...common import dtype as mstype


class PQC(PrimitiveWithInfer):
    r"""
    Evaluate a parameterized quantum circuit and calculate the gradient of each parameters.

    Inputs of this operation is generated by MindQuantum framework.

    Inputs:
        - **n_qubits** (int) - The qubit number of quantum simulator.
        - **encoder_params_names** (List[str]) - The parameters names of encoder circuit.
        - **ansatz_params_names** (List[str]) - The parameters names of ansatz circuit.
        - **gate_names** (List[str]) - The name of each gate.
        - **gate_matrix** (List[List[List[List[float]]]]) - Real part and image
            part of the matrix of quantum gate.
        - **gate_obj_qubits** (List[List[int]]) - Object qubits of each gate.
        - **gate_ctrl_qubits** (List[List[int]]) - Control qubits of each gate.
        - **gate_params_names** (List[List[str]]) - Parameter names of each gate.
        - **gate_coeff** (List[List[float]]) - Coefficient of each parameter of each gate.
        - **gate_requires_grad** (List[List[bool]]) - Whether to calculate gradient
            of parameters of gates.
        - **hams_pauli_coeff** (List[List[float]]) - Coefficient of pauli words.
        - **hams_pauli_word** (List[List[List[str]]]) - Pauli words.
        - **hams_pauli_qubit** (List[List[List[int]]]) - The qubit that pauli matrix act on.
        - **n_threads** (int) - Thread to evaluate input data.

    Outputs:
        - **expected_value** (Tensor) - The expected value of hamiltonian.
        - **g1** (Tensor) - Gradient of encode circuit parameters.
        - **g2** (Tensor) - Gradient of ansatz circuit parameters.

    Supported Platforms:
        ``CPU``
    """
    @prim_attr_register
    def __init__(self, n_qubits, encoder_params_names, ansatz_params_names,
                 gate_names, gate_matrix, gate_obj_qubits, gate_ctrl_qubits,
                 gate_params_names, gate_coeff, gate_requires_grad,
                 hams_pauli_coeff, hams_pauli_word, hams_pauli_qubit,
                 n_threads):
        self.init_prim_io_names(
            inputs=['encoder_data', 'ansatz_data'],
            outputs=['results', 'encoder_gradient', 'ansatz_gradient'])
        self.n_hams = len(hams_pauli_coeff)

    def check_shape_size(self, encoder_data, ansatz_data):
        if len(encoder_data) != 2:
            raise ValueError(f"For '{self.name}', the dimension of 'encoder_data' should be 2, "
                             f"but got {len(encoder_data)}.")
        if len(ansatz_data) != 1:
            raise ValueError(f"For '{self.name}', the dimension of 'ansatz_data' should be 1, "
                             f"but got {len(ansatz_data)}.")

    def infer_shape(self, encoder_data, ansatz_data):
        self.check_shape_size(encoder_data, ansatz_data)
        return [encoder_data[0], self.n_hams], [
            encoder_data[0], self.n_hams,
            len(self.encoder_params_names)
        ], [encoder_data[0], self.n_hams,
            len(self.ansatz_params_names)]

    def infer_dtype(self, encoder_data, ansatz_data):
        args = {'encoder_data': encoder_data, 'ansatz_data': ansatz_data}
        validator.check_tensors_dtypes_same_and_valid(args, mstype.float_type,
                                                      self.name)
        return encoder_data, encoder_data, encoder_data


class Evolution(PrimitiveWithInfer):
    r"""
    Inputs of this operation is generated by MindQuantum framework.

    Inputs:
        - **n_qubits** (int) - The qubit number of quantum simulator.
        - **param_names** (List[str]) - The parameters names.
        - **gate_names** (List[str]) - The name of each gate.
        - **gate_matrix** (List[List[List[List[float]]]]) - Real part and image
            part of the matrix of quantum gate.
        - **gate_obj_qubits** (List[List[int]]) - Object qubits of each gate.
        - **gate_ctrl_qubits** (List[List[int]]) - Control qubits of each gate.
        - **gate_params_names** (List[List[str]]) - Parameter names of each gate.
        - **gate_coeff** (List[List[float]]) - Coefficient of each parameter of each gate.
        - **gate_requires_grad** (List[List[bool]]) - Whether to calculate gradient
            of parameters of gates.
        - **hams_pauli_coeff** (List[List[float]]) - Coefficient of pauli words.
        - **hams_pauli_word** (List[List[List[str]]]) - Pauli words.
        - **hams_pauli_qubit** (List[List[List[int]]]) - The qubit that pauli matrix act on.

    Outputs:
        - **Quantum state** (Tensor) - The quantum state after evolution.

    Supported Platforms:
        ``CPU``
    """
    @prim_attr_register
    def __init__(self, n_qubits, param_names, gate_names, gate_matrix,
                 gate_obj_qubits, gate_ctrl_qubits, gate_params_names,
                 gate_coeff, gate_requires_grad, hams_pauli_coeff,
                 hams_pauli_word, hams_pauli_qubit):
        """Initialize Evolution"""
        self.init_prim_io_names(inputs=['param_data'], outputs=['state'])
        self.n_qubits = n_qubits

    def check_shape_size(self, param_data):
        if len(param_data) != 1:
            raise ValueError(f"For '{self.name}', the dimension of 'param_data' should be 1, "
                             f"but got {len(param_data)}.")

    def infer_shape(self, param_data):
        self.check_shape_size(param_data)
        return [1 << self.n_qubits, 2]

    def infer_dtype(self, param_data):
        args = {'param_data': param_data}
        validator.check_tensors_dtypes_same_and_valid(args, mstype.float_type,
                                                      self.name)
        return param_data