kshakeh.br

/* -------------------------------------------------------------------------- *
 *                                   OpenMM                                   *
 * -------------------------------------------------------------------------- *
 * This is part of the OpenMM molecular simulation toolkit originating from   *
 * Simbios, the NIH National Center for Physics-Based Simulation of           *
 * Biological Structures at Stanford, funded under the NIH Roadmap for        *
 * Medical Research, grant U54 GM072970. See https://simtk.org.               *
 *                                                                            *
 * Portions copyright (c) 2009 Stanford University and the Authors.           *
 * Authors: Mark Friedrichs, Mike Houston                                     *
 * Contributors:                                                              *
 *                                                                            *
 * This program is free software: you can redistribute it and/or modify       *
 * it under the terms of the GNU Lesser General Public License as published   *
 * by the Free Software Foundation, either version 3 of the License, or       *
 * (at your option) any later version.                                        *
 *                                                                            *
 * This program is distributed in the hope that it will be useful,            *
 * but WITHOUT ANY WARRANTY; without even the implied warranty of             *
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the              *
 * GNU Lesser General Public License for more details.                        *
 *                                                                            *
 * You should have received a copy of the GNU Lesser General Public License   *
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.      *
 * -------------------------------------------------------------------------- */

/*
 * This kernel constrains bonds involving hydrogen using an unrolled
 * shake implementation. We assume that a heavy atom can have at most 
 * 3 hydrogens attached. The hydrogens are expected to be identical
 *
 * For atoms that have ewer than 3 hydrogens
 *
 * Also we don't check the constraints!
 *  
 * We do a fixed number of iterations without checking for convergence
 * to avoid stalling
 * *************************************************************/

/* Fix for precision of terms first order in dt
 * To be used with corresponding update kernels
 * The input posqp is now changes to posq rather than
 * posq+changes
 * */
kernel void
kshakeh_fix1(
		float maxIterations,      //number of iterations
		float strwidth, //stream width of posq
		float inputTolerance,    //tolerance
		float4 atoms<>, //heavy0, h1, h2, h3
		float3 posq[][], //positions before update
		float3 posqp[][], //changes to positions
		float4 params<>, // (1/m0, mu/2, blensq, 1/m1 )
		out float3 cposq0<>, //constrained position for heavy atom
		out float3 cposq1<>, //ditto for h1
		out float3 cposq2<>, //ditto for h2
		out float3 cposq3<> //ditto for h3
		) {

	float2 ai, aj1, aj2, aj3; //2d indices, can be precalc.
	float i; //iteration count
	float3 xi, xj1, xj2, xj3; //coordinates
	float3 xpi, xpj1, xpj2, xpj3; //coordinates
	float3 rij1, rij2, rij3, rpij, dr;
	float rpsqij, rrpr, acor, tolerance, converged;
	float mask2, mask3;
	float diff;
	float rij1sq, rij2sq, rij3sq;
	float ld1, ld2, ld3;
	
	// ai.y  = floor( atoms.x / strwidth );
   ai.y  = round( (atoms.x - fmod( atoms.x, strwidth ))/strwidth );
	ai.x  = atoms.x - ai.y * strwidth;

	// aj1.y = floor( atoms.y / strwidth );
   aj1.y  = round( (atoms.y - fmod( atoms.y, strwidth ))/strwidth );
	aj1.x = atoms.y - aj1.y * strwidth;

	//If further hydrogens are absent,
	//just set to the coordinates of the first
	//so we don't have any junk memory accesses
	//or nans/infs in the calcs.	

	if( atoms.z > -0.5f ){
		aj2.y = round( (atoms.z - fmod( atoms.z, strwidth ))/strwidth );
		aj2.x = atoms.z - aj2.y * strwidth;
		mask2 = 1.0f;
	} else {
		aj2   = aj1; 
		mask2 = 0.0f;
	}

	if( atoms.w > -0.5f ){	
		aj3.y = round( (atoms.w - fmod( atoms.w, strwidth ))/strwidth );
		aj3.x = atoms.w - aj3.y * strwidth;
		mask3 = 1.0f;
	} else {
		aj3   = aj1;
		mask3 = 0.0f;
	}

	cposq0       = posq[ai];
	cposq1       = posq[aj1];
	cposq2       = posq[aj2];
	cposq3       = posq[aj3];
	
	xi           = cposq0;
	xj1          = cposq1;
	xj2          = cposq2;
	xj3          = cposq3;
	
	rij1         = xi - xj1;
	rij2         = xi - xj2;
	rij3         = xi - xj3;

	rij1sq       = dot( rij1, rij1 );
	rij2sq       = dot( rij2, rij2 );
	rij3sq       = dot( rij3, rij3 );

	ld1          = params.z - rij1sq;
	ld2          = params.z - rij2sq;
	ld3          = params.z - rij3sq;

	xpi          = posqp[ai];
	xpj1         = posqp[aj1];
	xpj2         = posqp[aj2];
	xpj3         = posqp[aj3];


	i           = 0.0f;
   converged   = 1.0f;
   tolerance   = 2.0f*inputTolerance;
	while( i < maxIterations && converged > 0.0f ){

		//First hydrogen

		rpij       = xpi - xpj1; //This is really rpij - rij
		rpsqij     = dot( rpij, rpij ); //This is really deltar ^ 2
		rrpr       = dot( rij1, rpij ); //This is r.deltar
      acor       = ( ld1 - 2.0f * rrpr - rpsqij ) * params.y / ( rrpr + rij1sq ) ; 
      diff       = abs( ld1 - 2.0f * rrpr - rpsqij ) / (params.z * tolerance );
      acor       = (diff < 1.0f) ? 0.0f : acor;
      converged  = abs( acor );

		dr         = rij1 * acor;
		xpi       += dr * params.x;
		xpj1      -= dr * params.w;

		//Second hydrogen

		rpij       = xpi - xpj2;
		rpsqij     = dot( rpij, rpij );
		rrpr       = dot( rij2, rpij );

      diff       = abs( ld2 - 2.0f * rrpr - rpsqij ) / (params.z * tolerance );
      acor       = mask2 * ( ld2 - 2.0f * rrpr - rpsqij ) * params.y / ( rrpr + rij2sq ) ; 
      acor       = (diff < 1.0f) ? 0.0f : acor;
      converged += abs( acor );

		dr         = rij2 * acor;
		xpi       += dr * params.x;
		xpj2      -= dr * params.w;

		//Third hydrogen

		rpij       = xpi - xpj3;
		rpsqij     = dot( rpij, rpij );
		rrpr       = dot( rij3, rpij );
      diff       = abs( ld3 - 2.0f * rrpr - rpsqij ) / (params.z * tolerance );
      acor       = mask3 * ( ld3 - 2.0f * rrpr - rpsqij ) * params.y / ( rrpr + rij3sq ) ;
      acor       = (diff < 1.0f) ? 0.0f : acor;
      converged += abs( acor );

		dr         = rij3 * acor;
		xpi       += dr * params.x;
		xpj3      -= dr * params.w;
		
		i += 1.0f;
	}

	//Output modified delta's

	cposq0 = xpi;
	cposq1 = xpj1;
	cposq2 = xpj2;
	cposq3 = xpj3;

}

kernel void kshakeh_update1_fix1( 
		float strwidth, //width of cposq streams
		float2 invmap<>,    //shakeh inverse map
		float3 posqp<>,    //deltas from sd2
		float3 cposq0[][], //constrained delta for heavy atom
		float3 cposq1[][], //ditto for h1
		float3 cposq2[][], //ditto for h2
		float3 cposq3[][], //ditto for h3
		out float3 oposq<> //updated deltas
		){

	float2 atom;
	
   atom.y  = round( (invmap.x - fmod( invmap.x, strwidth ))/strwidth );
	atom.x  = invmap.x - atom.y * strwidth;
	
	if( invmap.y < 0 ){
		oposq  = posqp;
	} else if( invmap.y < 0.5f ){
		oposq  = cposq0[ atom ];
	} else if( invmap.y < 1.5f ){
		oposq  = cposq1[ atom ];
	} else if( invmap.y < 2.5f ){
		oposq  = cposq2[ atom ];
	} else if( invmap.y < 3.5f ){
		oposq  = cposq3[ atom ];
   }

}

kernel void kshakeh_update2_fix1( 
		float strwidth, //width of cposq streams
		float2 invmap<>,    //shakeh inverse map
		float3 posq<>,     //old positions
		float3 posqp<>,    //deltas from sd2
		float3 cposq0[][], //constrained delta for heavy atom
		float3 cposq1[][], //ditto for h1
		float3 cposq2[][], //ditto for h2
		float3 cposq3[][], //ditto for h3
		out float3 oposq<> //updated deltas
		){

	float2 atom;
	
   atom.y  = round( (invmap.x - fmod( invmap.x, strwidth ))/strwidth );
	atom.x  = invmap.x - atom.y * strwidth;
	
	if ( invmap.y < 0 ){
		oposq = posqp;
	} else if ( invmap.y < 0.5f ){ 
		oposq = cposq0[ atom ];
	} else if ( invmap.y < 1.5f){
		oposq = cposq1[ atom ];
	} else if ( invmap.y < 2.5f ){
		oposq = cposq2[ atom ];
	} else if ( invmap.y < 3.5f ){
		oposq = cposq3[ atom ];
   }

	oposq += posq;
}