Amazon SageMaker で量子系をシミュレートする

git clone https://github.com/emkessler/sagemaker.git

import subprocess
print subprocess.check_output(['./build_and_push.sh', 'superradiance'])

from scipy import sparse
import numpy as np
import os
import pickle
import matplotlib.pyplot as plt

import boto3
from sagemaker import get_execution_role
import sagemaker as sage

steps = 500                  # Number of simulated time steps
N = 8                        # Number of nuclear spins
dim = 2**(N+1)               # Dimension of Hilbert space of N nuclear and 1 electronic spin

# Define Nuclear and Electron spin states as density matrices
rhoI = sparse.csr_matrix(([1.],([0],[0])),shape =(int(dim/2),int(dim/2))) # All carbon atoms magnetically excited
rhoS = sparse.csr_matrix(([1.],([1],[1])),shape =(2,2))                   # Electron magnetically excited
sparse_rho =sparse.kron(rhoS, rhoI)                                       # Build the system density matrix

tempfile = '/tmp/tmp.pckl'
pickle.dump(sparse_rho, open(tempfile, "wb"))

# Upload serialized initial state to S3
resource = boto3.resource('s3') 
my_bucket = resource.Bucket('sagemaker-kessle31')                         #subsitute this for your s3 bucket name. 
my_bucket.upload_file(tempfile, Key='superradiance/initial_state/init.pckl')

# Clean up temporary files
os.remove(tempfile) 

role = get_execution_role()
sess = sage.Session()

# Specify the name of the ECS repo that contains your docker image. 
# By default the image with tag 'latest' in the repo will be utilized.
imagename = 'superradiance'  

account = sess.boto_session.client('sts').get_caller_identity()['Account']
region = sess.boto_session.region_name
image = '{}.dkr.ecr.{}.amazonaws.com/{}'.format(account, region, imagename)

# Define the training job
superradiance = sage.estimator.Estimator(image,
                       role, 1, 'ml.c4.2xlarge',
                       output_path="s3://sagemaker-kessle31/superradiance/output",
                       sagemaker_session=sess)

# Parameters can be passed to the simulation as a dictionary
superradiance.hyperparam_dict = {'N': N, 'steps': steps}

# Pass the location of the training data (see above) and start the job
superradiance.fit("s3://sagemaker-kessle31/superradiance/initial_state")
INFO:sagemaker:Creating training-job with name: superradiance-2018-02-22-16-31-48-732

import tarfile

# get the results from S3
results = my_bucket.Object('superradiance/output/{}/output/model.tar.gz'.format(superradiance.latest_training_job.name))
tempfile = '/tmp/model.tar.gz'   
results.download_file('/tmp/model.tar.gz')    

# unzip the results
tar = tarfile.open(tempfile, "r:gz")
tar.extractall(path = '/tmp')

# load the results into the notebook
out = pickle.load(open('/tmp/out.pckl', "rb"))
intensity = out['intensity']
I_ind = out['I_ind']

# clean-up temporary files
os.remove(tempfile) 

# Vizualize the results
plt.scatter(range(steps-1),intensity/I_ind, label = 'Simulated System')
plt.scatter(range(steps-1), np.exp(-I_ind/N * np.arange(steps-1)), label = 'Classical System (benchmark)')
plt.xlabel('Time [a.u.]')
plt.ylabel('Intensity of emitted radiation [normalized]')
plt.legend()

plt.show()

!cat container/superradiance/train
#!/usr/bin/env python

# A sample program that simulates a simple quantum system exhibiting superradiance

from __future__ import print_function
import os
import json
import pickle
import sys
import traceback
import tensorflow as tf
import numpy as np
from utils import kronecker_product, SparseIndices               
from scipy import sparse                                          
import pandas as pd                                               

# These are the paths to where SageMaker mounts interesting things in your container.

prefix = '/opt/ml/'

input_path = prefix + 'input/data'

output_path = os.path.join(prefix, 'output')                       # Failure output & error messages should be written here
model_path = os.path.join(prefix, 'model')                         # All results should be written here 
param_path = os.path.join(prefix, 
                          'input/config/hyperparameters.json')     # Passed parameters can be retrieved here

# This algorithm has a single channel of input data called 'training'. Since we run in
# File mode, the input files are copied to the directory specified here.
channel_name='training'
training_path = os.path.join(input_path, channel_name)              # The initial state is found here 



# The function to execute the training.
def train():
    print('Starting the training.')
    try:        
        ##################################################################################
        # Retrieve parameters
        ##################################################################################
        # Parameters that were defined in the hyperparameter dict in the notebook 
        # (in our example: superradiance.hyperparam_dict = {'N': N, 'steps': steps})  
        # can be retrieved from the param_path defined above. 
        # Note that the parameters are stored as a JSON string, so it is required to 
        # define the correct dtype upon reading (here: int).
        
        datatype = tf.complex128
        
        with open(param_path, 'r') as tc: 
             trainingParams = json.load(tc)
        N = int(trainingParams.get('N', None))          # Number of simulated nuclear spins as defined in notebook
        steps = int(trainingParams.get('steps', None))  # Number of simulated time steps

        dim = 2**(N+1)                                  # Dimension of Hilbert space of N nuclear and 1 electronic spin
        
        #################################################################################
        # Define some other parameters
        ##################################################################################
        
        Amean = 1
        gamma = 1
        #g = np.random.randn(N) 
        g = np.ones(N)
        g = g/(np.sqrt((g**2).sum()))                    # Coupling strengts
        A = Amean/g.sum()
        I_ind = np.real(A**2/(gamma+1j*Amean))*2/N       # Classical emission intensity of N independent emitters (benchmark)        
        
        ##################################################################################
        # Define TF constants
        ##################################################################################
        
        gammaT = tf.constant(gamma, dtype=datatype)
        AT = tf.constant(A, dtype=datatype)
        eps = tf.placeholder(dtype=datatype, shape=())

        ##################################################################################
        # Define 'Pauli Matrices'
        ##################################################################################
        # The most elemental quantum mechanical spins system is a spin-1/2. It can be 
        # represented through three simple matrices, called Pauli Matrices. For more 
        # information, there are many great online resources, e.g., 
        # http://galileo.phys.virginia.edu/~pf7a/msm3.pdf

        sigmaz = tf.constant([[1/2, 0], [0, -1/2]], dtype= datatype)
        splus  = tf.constant([[0, 1], [0, 0]], dtype = datatype)
        sminus = tf.constant([[0, 0], [1, 0]], dtype = datatype)

        ##################################################################################
        # Electron Spin operators 
        ##################################################################################
        # States of a quantum system live in a vector space with special properties, 
        # called the Hilbert space. 
        # The Hilbert space of a composite system of multiple spins can be constructed by
        # taking the tensor product of the individual 1-particle spaces. 
        # Using this property, it can be shown that we can construct the multi-spin 
        # operators by taking the the Kronecker product (https://en.wikipedia.org/wiki/Kronecker_product)
        # of the individual spin operators with the identity matrix of appropriate dimension 
        # (see, section 3.2 of http://galileo.phys.virginia.edu/~pf7a/msm3.pdf)
        # Let's see how that looks like for our electron operator:

        Sm = kronecker_product(sminus, tf.eye(int(dim/2), dtype = datatype))
        Sp = tf.transpose(Sm)
        Sz = kronecker_product(sigmaz, tf.eye(int(dim/2), dtype = datatype))

        ##################################################################################
        # Nuclear Spin operators 
        ##################################################################################
        # Similarly to the case of the electron operator, we create one spin operator 
        # for every nuclear spin we simulate, and take the sum over all nuclear operators 
        # (because the electron interacts with all of them)

        Am = tf.zeros(int(dim/2), dtype = datatype)
        Iz = tf.zeros(int(dim/2), dtype = datatype)
        for p in range(N): 
            Am = Am + g[p]*kronecker_product(kronecker_product(tf.eye(2**p, dtype = datatype),sminus)
                                           ,tf.eye(2**(N-1-p), dtype = datatype))
            Iz = Iz +      kronecker_product(kronecker_product(tf.eye(2**p, dtype = datatype),sigmaz)
                                           ,tf.eye(2**(N-1-p), dtype = datatype))

        Am = kronecker_product(tf.eye(2, dtype = datatype),Am)
        Ap = tf.transpose(Am)
        Iz = kronecker_product(tf.eye(2, dtype = datatype),Iz)
            
        ##################################################################################
        # Load and construct Initial State
        ##################################################################################
        # You can pass on larger data structures to your program by placing them on S3 and 
        # and point to the location at execution time.
        # (in our example: superradiance.fit("s3://sagemaker-kessle31/superradiance/initial_state")
        # It can be retrieved from the 'training_path' location defined above.           
            
        sparse_rho = pickle.load(open(os.path.join(training_path, 'init.pckl'), "rb"))

        rho_ind = SparseIndices(sparse_rho)
        rho2 = tf.sparse_to_dense(rho_ind['indices'], 
                                  rho_ind['dense_shape'], 
                                  rho_ind['values'])         # Sparse_to_dense not implemented for complex at time of this 
                                                             # notebook => Construct Real an Imag separately
        rho1 = tf.complex(rho2,tf.zeros(rho_ind['dense_shape'], dtype=tf.float64))
        rho = tf.Variable(rho1, dtype=datatype)  

        ##################################################################################
        # Define Operations
        ##################################################################################
        # The Hamiltonian and Lagrangian are operators codifying the dynamics and interactions 
        # of a Quantum System.
        # The system we are simulating is a simple central spin system represented by the 
        # following Hamiltonian/Lagrangian:
        # (https://journals.aps.org/prl/abstract/10.1103/PhysRevLett.104.143601)  
        
        H = 0.5 * AT * (tf.matmul(Sp,Am) + tf.matmul(Sm,Ap))  # Define the Hamiltonian of the system
        
        IL = gammaT*(tf.matmul(tf.matmul(Sm,rho),Sp) 
                     - 0.5 * (tf.matmul(tf.matmul(Sp,Sm),rho) 
                              + tf.matmul(rho, tf.matmul(Sp,Sm))))\
              - 1j * (tf.matmul(H,rho) - tf.matmul(rho, H))   # Define the Lagrangian of the system
            
        rho_ = rho + eps * IL                                 # Evolve the system by one time step
        
        pol = tf.real(tf.trace(tf.matmul(rho, Iz)))           # Compute the polarization of the nuclear system,

        global_step = tf.train.get_or_create_global_step()
        increment_global_step_op = tf.assign(global_step, global_step+1)
        
        step = tf.group(rho.assign(rho_),increment_global_step_op) 
        
        ##################################################################################
        # Run the Simulation     
        ##################################################################################
        
        polarization = np.zeros(steps)        
        init = tf.global_variables_initializer()
        with tf.Session() as sess:
            sess.run(init)            
            for i in range(steps):
                _ , polarization[i] = sess.run([step, pol], feed_dict={eps:0.3})
                
        ##################################################################################
        # Compute output and dump to pickle                       
        ##################################################################################                       
        # Everything written into the 'model_path' location defined above will be
        # returned by SageMaker to the S3 location defined in the notebook 
        # (in our example: output_path="s3://sagemaker-kessle31/superradiance/output")
        # Note that the files will be returned as a tar.gz archive       
        
        intensity = -np.diff(polarization)                                  
        out= {'intensity': intensity, 'I_ind': I_ind}
        pickle.dump(out, open(os.path.join(model_path, 'out.pckl'), "wb"))
        
        print('Training complete.')
        
    except Exception as e:
        # Write out an error file. This will be returned as the failureReason in the
        # DescribeTrainingJob result.
        trc = traceback.format_exc()
        with open(os.path.join(output_path, 'failure'), 'w') as s:
            s.write('Exception during training: ' + str(e) + '\n' + trc)
        # Printing this causes the exception to be in the training job logs, as well.
        print('Exception during training: ' + str(e) + '\n' + trc, file=sys.stderr)
        # A non-zero exit code causes the training job to be marked as Failed.
        sys.exit(255)

if __name__ == '__main__':
    train()

    # A zero exit code causes the job to be marked a Succeeded.
    sys.exit(0)