How to use the molgrid.Transform function in molgrid

def test_random_transform():
  from molgrid import Transform
  molgrid.set_random_seed(0)
  c1 = molgrid.float3(0,0,0);
  c2 = molgrid.float3(0,0,1);

  t1 = Transform(c1, 4.0, True)
  t2 = Transform(c2, 4.0, True)

  nrt1 = Transform (c1)
  nrt2 = Transform (c2)

  t = Transform()

  molgrid.set_random_seed(0) # reset, should get same sample
  t3 = Transform(c1, 4.0, True);

  neqQ(t1.get_quaternion(),t2.get_quaternion());
  neqQ(t1.get_quaternion(),nrt1.get_quaternion());
  eqQ(t1.get_quaternion(),t3.get_quaternion());
  eqQ(nrt1.get_quaternion(),nrt2.get_quaternion());

  assert tup(t1.get_translation()) != tup(t2.get_translation())
  assert tup(t1.get_translation()) != tup(nrt1.get_translation())
  assert tup(t1.get_translation()) == tup(t3.get_translation())
  assert tup(nrt1.get_translation()) == tup(nrt2.get_translation())

  assert tup(c1) == tup(t1.get_rotation_center())
  assert tup(c2) == tup(t2.get_rotation_center())

g1 = molgrid.GridMaker(resolution=.1,dimension=6.0)
    c = np.array([[1.0,0,0]],np.float32)
    t = np.array([0],np.float32)
    r = np.array([2.0],np.float32)
    coords = molgrid.CoordinateSet(c,t,r,1)
    shape = g1.grid_dimensions(1)
    
    #make diff with gradient in center
    diff = molgrid.MGrid4f(*shape)
    diff[0,30,30,30] = 1.0  
    
    cpuatoms = molgrid.MGrid2f(1,3)
    gpuatoms = molgrid.MGrid2f(1,3)
    
    #apply random rotation
    T = molgrid.Transform((0,0,0), 0, True)
    T.forward(coords, coords);
    
    g1.backward((0,0,0),coords,diff.cpu(), cpuatoms.cpu())
    g1.backward((0,0,0),coords,diff.gpu(), gpuatoms.gpu())

    T.backward(cpuatoms.cpu(), cpuatoms.cpu(), False)
    T.backward(gpuatoms.gpu(), gpuatoms.gpu(), False)
    
    print(cpuatoms.tonumpy(), gpuatoms.tonumpy())
    # results should be ~ -.6, 0, 0
    np.testing.assert_allclose(cpuatoms.tonumpy(), gpuatoms.tonumpy(), atol=1e-5)
    np.testing.assert_allclose(cpuatoms.tonumpy().flatten(), [-0.60653067,0,0], atol=1e-5)

def test_random_transform():
  from molgrid import Transform
  molgrid.set_random_seed(0)
  c1 = molgrid.float3(0,0,0);
  c2 = molgrid.float3(0,0,1);

  t1 = Transform(c1, 4.0, True)
  t2 = Transform(c2, 4.0, True)

  nrt1 = Transform (c1)
  nrt2 = Transform (c2)

  t = Transform()

  molgrid.set_random_seed(0) # reset, should get same sample
  t3 = Transform(c1, 4.0, True);

  neqQ(t1.get_quaternion(),t2.get_quaternion());
  neqQ(t1.get_quaternion(),nrt1.get_quaternion());
  eqQ(t1.get_quaternion(),t3.get_quaternion());
  eqQ(nrt1.get_quaternion(),nrt2.get_quaternion());

  assert tup(t1.get_translation()) != tup(t2.get_translation())
  assert tup(t1.get_translation()) != tup(nrt1.get_translation())
  assert tup(t1.get_translation()) == tup(t3.get_translation())
  assert tup(nrt1.get_translation()) == tup(nrt2.get_translation())

def test_random_transform():
  from molgrid import Transform
  molgrid.set_random_seed(0)
  c1 = molgrid.float3(0,0,0);
  c2 = molgrid.float3(0,0,1);

  t1 = Transform(c1, 4.0, True)
  t2 = Transform(c2, 4.0, True)

  nrt1 = Transform (c1)
  nrt2 = Transform (c2)

  t = Transform()

  molgrid.set_random_seed(0) # reset, should get same sample
  t3 = Transform(c1, 4.0, True);

  neqQ(t1.get_quaternion(),t2.get_quaternion());
  neqQ(t1.get_quaternion(),nrt1.get_quaternion());
  eqQ(t1.get_quaternion(),t3.get_quaternion());
  eqQ(nrt1.get_quaternion(),nrt2.get_quaternion());

  assert tup(t1.get_translation()) != tup(t2.get_translation())
  assert tup(t1.get_translation()) != tup(nrt1.get_translation())
  assert tup(t1.get_translation()) == tup(t3.get_translation())
  assert tup(nrt1.get_translation()) == tup(nrt2.get_translation())

  assert tup(c1) == tup(t1.get_rotation_center())
  assert tup(c2) == tup(t2.get_rotation_center())
  assert tup(c1) == tup(nrt1.get_rotation_center())

  assert tup(c1) == tup(t.get_rotation_center())

for i in range(1,batch_size):
        sqsum = np.square(b[0].coord_sets[1].coords.tonumpy() - b[i].coord_sets[1].coords.tonumpy()).sum()
        assert sqsum > 0
        
    #now with duplicates
    e = molgrid.ExampleProvider(data_root=datadir+"/structs",num_copies=batch_size)
    e.populate(fname)
    b = e.next_batch(batch_size)
    for i in range(1,batch_size):
        sqsum = np.square(b[0].coord_sets[1].coords.tonumpy() - b[i].coord_sets[1].coords.tonumpy()).sum()
        assert sqsum == 0    
        
    #transforming one of the duplicates should not effect the others
    orig = b[0].coord_sets[1].coords.tonumpy()
    lig = b[1].coord_sets[1]
    t = molgrid.Transform(lig.center(), 2,random_rotation=True)
    t.forward(lig,lig)
    new0 = b[0].coord_sets[1].coords.tonumpy()
    new1 = b[1].coord_sets[1].coords.tonumpy()
    
    np.testing.assert_allclose(orig,new0)
    sqsum = np.square(new1-orig).sum()
    assert sqsum > 0

def test_random_transform():
  from molgrid import Transform
  molgrid.set_random_seed(0)
  c1 = molgrid.float3(0,0,0);
  c2 = molgrid.float3(0,0,1);

  t1 = Transform(c1, 4.0, True)
  t2 = Transform(c2, 4.0, True)

  nrt1 = Transform (c1)
  nrt2 = Transform (c2)

  t = Transform()

  molgrid.set_random_seed(0) # reset, should get same sample
  t3 = Transform(c1, 4.0, True);

  neqQ(t1.get_quaternion(),t2.get_quaternion());
  neqQ(t1.get_quaternion(),nrt1.get_quaternion());
  eqQ(t1.get_quaternion(),t3.get_quaternion());
  eqQ(nrt1.get_quaternion(),nrt2.get_quaternion());

  assert tup(t1.get_translation()) != tup(t2.get_translation())
  assert tup(t1.get_translation()) != tup(nrt1.get_translation())
  assert tup(t1.get_translation()) == tup(t3.get_translation())
  assert tup(nrt1.get_translation()) == tup(nrt2.get_translation())