#!/usr/bin/env python
# -*- coding: utf-8 -*-

import importlib

import ctypes as ct
import pysixtracklib as pyst
from pysixtracklib import stcommon as st
import numpy as np

pycuda_spec = importlib.util.find_spec('pycuda')

if pycuda_spec is not None:
    import pycuda
    import pycuda.gpuarray
    import pycuda.driver
    import pycuda.autoinit

if numpy_spec is not None:


if __name__ == '__main__':
    if not pyst.supports('cuda'):
        raise SystemExit("Example requires cuda support in pysixtracklib")

    if pycuda_spec is None:
        raise SystemExit("Example requires pycuda installation")

    num_particles = 42
    partset = pyst.ParticlesSet()
    particles = partset.Particles(num_particles=num_particles)

    cmp_partset = pyst.ParticlesSet()
    cmp_particles = cmp_partset.Particles(num_particles=num_particles)

    elements = pyst.Elements()
    elements.Drift(length=1.2)
    elements.Multipole(knl=[0, 0.001])

    track_job = pyst.CudaTrackJob(elements, partset)

    if not track_job.has_particles_addresses and \
            track_job.can_fetch_particle_addresses:
        track_job.fetch_particle_addresses()

    pset_index = 0
    ptr_particles_addr = track_job.get_particle_addresses(pset_index)
    particles_addr = ptr_particles_addr.contents

    print("Particle structure data on the device:")
    print("num_particles = {0:8d}".format(particles_addr.num_particles))
    print("x     begin at = {0:16x}".format(particles_addr.x))
    print("y     begin at = {0:16x}".format(particles_addr.y))
    print("px    begin at = {0:16x}".format(particles_addr.px))
    print("py    begin at = {0:16x}".format(particles_addr.py))
    print("zeta  begin at = {0:16x}".format(particles_addr.zeta))
    print("delta begin at = {0:16x}".format(particles_addr.delta))

    x = pycuda.gpuarray.GPUArray(
        particles_addr.num_particles, float, gpudata=particles_addr.x)

    x = np.linspace(0.0, float(num_particles - 1), num=num_particles,
                    dtype=np.float64)

    cmp_particles.x = np.linspace(
        0.0, float(num_particles - 1), num=num_particles, dtype=np.float64)

    assert pyst.compareParticlesDifference(
        cmp_particles, particles, abs_treshold=2e-14) != 0

    track_job.collectParticles()

    assert pyst.compareParticlesDifference(
        cmp_particles, particles, abs_treshold=2e-14) == 0

@@ -57,16 +57,42 @@
 portlib%0A
+from importlib import util
 %0Aimport 
@@ -185,45 +185,65 @@
 ort 
-numpy as np%0A%0Apycuda_spec = importlib.
+pdb%0A%0Apycuda_spec = util.find_spec('pycuda')%0Anumpy_spec = 
 util
@@ -254,22 +254,21 @@
 d_spec('
-pycuda
+numpy
 ')%0A%0Aif p
@@ -416,16 +416,38 @@
 t None:%0A
+    import numpy as np
 %0A%0Aif __n
@@ -667,24 +667,116 @@
 allation%22)%0A%0A
+    if numpy_spec is None:%0A        raise SystemExit(%22Example requires numpy installation%22)%0A%0A
     num_part
@@ -1096,16 +1096,36 @@
 .001%5D)%0A%0A
+    pdb.set_trace()%0A
     trac
@@ -1199,25 +1199,24 @@
 has_particle
-s
 _addresses a
@@ -1534,16 +1534,17 @@
 rticles 
+ 
 = %7B0:8d%7D
@@ -1971,16 +1971,37 @@
 elta))%0A%0A
+    pdb.set_trace()%0A%0A
     x = 

update cuda fetch addresses plus pycuda interfacing example