kshakeh.br 8.05 KB
Newer Older
Mark Friedrichs's avatar
Mark Friedrichs committed
1

2
3
4
5
6
7
8
9
/* -------------------------------------------------------------------------- *
 *                                   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.               *
 *                                                                            *
10
11
 * Portions copyright (c) 2009 Stanford University and the Authors.           *
 * Authors: Mark Friedrichs, Mike Houston                                     *
12
13
 * Contributors:                                                              *
 *                                                                            *
14
15
16
17
 * 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.                                        *
18
 *                                                                            *
19
20
21
22
 * 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.                        *
23
 *                                                                            *
24
25
 * 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/>.      *
26
 * -------------------------------------------------------------------------- */
Mark Friedrichs's avatar
Mark Friedrichs committed
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47

/*
 * 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(
Mark Friedrichs's avatar
Mods  
Mark Friedrichs committed
48
		float maxIterations,      //number of iterations
Mark Friedrichs's avatar
Mark Friedrichs committed
49
		float strwidth, //stream width of posq
Mark Friedrichs's avatar
Mods  
Mark Friedrichs committed
50
		float inputTolerance,    //tolerance
Mark Friedrichs's avatar
Mark Friedrichs committed
51
		float4 atoms<>, //heavy0, h1, h2, h3
52
53
		float3 posq[][], //positions before update
		float3 posqp[][], //changes to positions
Mark Friedrichs's avatar
Mods  
Mark Friedrichs committed
54
		float4 params<>, // (1/m0, mu/2, blensq, 1/m1 )
55
56
57
58
		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
Mark Friedrichs's avatar
Mark Friedrichs committed
59
		) {
Mark Friedrichs's avatar
Mods  
Mark Friedrichs committed
60

Mark Friedrichs's avatar
Mark Friedrichs committed
61
62
63
64
65
	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;
Mark Friedrichs's avatar
Mods  
Mark Friedrichs committed
66
	float rpsqij, rrpr, acor, tolerance, converged;
Mark Friedrichs's avatar
Mark Friedrichs committed
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
	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.	
Mark Friedrichs's avatar
Mods  
Mark Friedrichs committed
84
85

	if( atoms.z > -0.5f ){
Mark Friedrichs's avatar
Mark Friedrichs committed
86
87
88
		aj2.y = round( (atoms.z - fmod( atoms.z, strwidth ))/strwidth );
		aj2.x = atoms.z - aj2.y * strwidth;
		mask2 = 1.0f;
Mark Friedrichs's avatar
Mods  
Mark Friedrichs committed
89
	} else {
Mark Friedrichs's avatar
Mods  
Mark Friedrichs committed
90
		aj2   = aj1; 
Mark Friedrichs's avatar
Mark Friedrichs committed
91
92
93
		mask2 = 0.0f;
	}

Mark Friedrichs's avatar
Mods  
Mark Friedrichs committed
94
	if( atoms.w > -0.5f ){	
Mark Friedrichs's avatar
Mark Friedrichs committed
95
96
97
		aj3.y = round( (atoms.w - fmod( atoms.w, strwidth ))/strwidth );
		aj3.x = atoms.w - aj3.y * strwidth;
		mask3 = 1.0f;
Mark Friedrichs's avatar
Mods  
Mark Friedrichs committed
98
	} else {
Mark Friedrichs's avatar
Mods  
Mark Friedrichs committed
99
		aj3   = aj1;
Mark Friedrichs's avatar
Mark Friedrichs committed
100
101
102
		mask3 = 0.0f;
	}

Mark Friedrichs's avatar
Mods  
Mark Friedrichs committed
103
104
105
106
	cposq0       = posq[ai];
	cposq1       = posq[aj1];
	cposq2       = posq[aj2];
	cposq3       = posq[aj3];
Mark Friedrichs's avatar
Mark Friedrichs committed
107
	
Mark Friedrichs's avatar
Mods  
Mark Friedrichs committed
108
109
110
111
	xi           = cposq0;
	xj1          = cposq1;
	xj2          = cposq2;
	xj3          = cposq3;
Mark Friedrichs's avatar
Mark Friedrichs committed
112
	
Mark Friedrichs's avatar
Mods  
Mark Friedrichs committed
113
114
115
	rij1         = xi - xj1;
	rij2         = xi - xj2;
	rij3         = xi - xj3;
Mark Friedrichs's avatar
Mark Friedrichs committed
116

Mark Friedrichs's avatar
Mods  
Mark Friedrichs committed
117
118
119
	rij1sq       = dot( rij1, rij1 );
	rij2sq       = dot( rij2, rij2 );
	rij3sq       = dot( rij3, rij3 );
Mark Friedrichs's avatar
Mark Friedrichs committed
120

Mark Friedrichs's avatar
Mods  
Mark Friedrichs committed
121
122
123
	ld1          = params.z - rij1sq;
	ld2          = params.z - rij2sq;
	ld3          = params.z - rij3sq;
Mark Friedrichs's avatar
Mark Friedrichs committed
124

Mark Friedrichs's avatar
Mods  
Mark Friedrichs committed
125
126
127
128
	xpi          = posqp[ai];
	xpj1         = posqp[aj1];
	xpj2         = posqp[aj2];
	xpj3         = posqp[aj3];
Mark Friedrichs's avatar
Mark Friedrichs committed
129
130


Mark Friedrichs's avatar
Mods  
Mark Friedrichs committed
131
132
	i           = 0.0f;
   converged   = 1.0f;
Mark Friedrichs's avatar
Mods  
Mark Friedrichs committed
133
   tolerance   = 2.0f*inputTolerance;
Mark Friedrichs's avatar
Mods  
Mark Friedrichs committed
134
	while( i < maxIterations && converged > 0.0f ){
Mark Friedrichs's avatar
Mark Friedrichs committed
135
136
137

		//First hydrogen

Mark Friedrichs's avatar
Mods  
Mark Friedrichs committed
138
139
140
		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
Mark Friedrichs's avatar
Mods  
Mark Friedrichs committed
141
142
      acor       = ( ld1 - 2.0f * rrpr - rpsqij ) * params.y / ( rrpr + rij1sq ) ; 
      diff       = abs( ld1 - 2.0f * rrpr - rpsqij ) / (params.z * tolerance );
Mark Friedrichs's avatar
Mods  
Mark Friedrichs committed
143
      acor       = (diff < 1.0f) ? 0.0f : acor;
Mark Friedrichs's avatar
Mods  
Mark Friedrichs committed
144
      converged  = abs( acor );
Mark Friedrichs's avatar
Mark Friedrichs committed
145

Mark Friedrichs's avatar
Mods  
Mark Friedrichs committed
146
147
		dr         = rij1 * acor;
		xpi       += dr * params.x;
Mark Friedrichs's avatar
Mods  
Mark Friedrichs committed
148
		xpj1      -= dr * params.w;
Mark Friedrichs's avatar
Mark Friedrichs committed
149

Mark Friedrichs's avatar
Mods  
Mark Friedrichs committed
150
		//Second hydrogen
Mark Friedrichs's avatar
Mark Friedrichs committed
151

Mark Friedrichs's avatar
Mods  
Mark Friedrichs committed
152
153
154
		rpij       = xpi - xpj2;
		rpsqij     = dot( rpij, rpij );
		rrpr       = dot( rij2, rpij );
Mark Friedrichs's avatar
Mark Friedrichs committed
155

Mark Friedrichs's avatar
Mods  
Mark Friedrichs committed
156
157
158
      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;
Mark Friedrichs's avatar
Mods  
Mark Friedrichs committed
159
      converged += abs( acor );
Mark Friedrichs's avatar
Mark Friedrichs committed
160

Mark Friedrichs's avatar
Mods  
Mark Friedrichs committed
161
162
		dr         = rij2 * acor;
		xpi       += dr * params.x;
Mark Friedrichs's avatar
Mods  
Mark Friedrichs committed
163
		xpj2      -= dr * params.w;
Mark Friedrichs's avatar
Mark Friedrichs committed
164
165

		//Third hydrogen
Mark Friedrichs's avatar
Mods  
Mark Friedrichs committed
166
167
168
169
170
171
172

		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;
Mark Friedrichs's avatar
Mods  
Mark Friedrichs committed
173
      converged += abs( acor );
Mark Friedrichs's avatar
Mods  
Mark Friedrichs committed
174
175
176

		dr         = rij3 * acor;
		xpi       += dr * params.x;
Mark Friedrichs's avatar
Mods  
Mark Friedrichs committed
177
		xpj3      -= dr * params.w;
Mark Friedrichs's avatar
Mark Friedrichs committed
178
179
180
181
182
		
		i += 1.0f;
	}

	//Output modified delta's
Mark Friedrichs's avatar
Mods  
Mark Friedrichs committed
183

184
185
186
187
	cposq0 = xpi;
	cposq1 = xpj1;
	cposq2 = xpj2;
	cposq3 = xpj3;
Mark Friedrichs's avatar
Mark Friedrichs committed
188
189
190

}

Mark Friedrichs's avatar
Mods  
Mark Friedrichs committed
191
kernel void kshakeh_update1_fix1( 
Mark Friedrichs's avatar
Mark Friedrichs committed
192
193
		float strwidth, //width of cposq streams
		float2 invmap<>,    //shakeh inverse map
194
195
196
197
198
199
		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
Mark Friedrichs's avatar
Mark Friedrichs committed
200
		){
Mark Friedrichs's avatar
Mods  
Mark Friedrichs committed
201

Mark Friedrichs's avatar
Mark Friedrichs committed
202
203
204
	float2 atom;
	
   atom.y  = round( (invmap.x - fmod( invmap.x, strwidth ))/strwidth );
Mark Friedrichs's avatar
Mods  
Mark Friedrichs committed
205
	atom.x  = invmap.x - atom.y * strwidth;
Mark Friedrichs's avatar
Mark Friedrichs committed
206
	
Mark Friedrichs's avatar
Mods  
Mark Friedrichs committed
207
208
209
210
211
212
213
214
215
216
	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 ];
Mark Friedrichs's avatar
Mods  
Mark Friedrichs committed
217
   }
Mark Friedrichs's avatar
Mark Friedrichs committed
218
219
220

}

Mark Friedrichs's avatar
Mods  
Mark Friedrichs committed
221
kernel void kshakeh_update2_fix1( 
Mark Friedrichs's avatar
Mark Friedrichs committed
222
223
		float strwidth, //width of cposq streams
		float2 invmap<>,    //shakeh inverse map
224
225
226
227
228
229
230
		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
Mark Friedrichs's avatar
Mark Friedrichs committed
231
		){
Mark Friedrichs's avatar
Mods  
Mark Friedrichs committed
232

Mark Friedrichs's avatar
Mark Friedrichs committed
233
234
235
	float2 atom;
	
   atom.y  = round( (invmap.x - fmod( invmap.x, strwidth ))/strwidth );
Mark Friedrichs's avatar
Mods  
Mark Friedrichs committed
236
	atom.x  = invmap.x - atom.y * strwidth;
Mark Friedrichs's avatar
Mark Friedrichs committed
237
238
239
	
	if ( invmap.y < 0 ){
		oposq = posqp;
Mark Friedrichs's avatar
Mods  
Mark Friedrichs committed
240
241
242
243
	} else if ( invmap.y < 0.5f ){ 
		oposq = cposq0[ atom ];
	} else if ( invmap.y < 1.5f){
		oposq = cposq1[ atom ];
Mark Friedrichs's avatar
Mods  
Mark Friedrichs committed
244
	} else if ( invmap.y < 2.5f ){
Mark Friedrichs's avatar
Mods  
Mark Friedrichs committed
245
		oposq = cposq2[ atom ];
Mark Friedrichs's avatar
Mods  
Mark Friedrichs committed
246
	} else if ( invmap.y < 3.5f ){
Mark Friedrichs's avatar
Mods  
Mark Friedrichs committed
247
		oposq = cposq3[ atom ];
Mark Friedrichs's avatar
Mods  
Mark Friedrichs committed
248
   }
Mark Friedrichs's avatar
Mark Friedrichs committed
249

Mark Friedrichs's avatar
Mods  
Mark Friedrichs committed
250
	oposq += posq;
Mark Friedrichs's avatar
Mark Friedrichs committed
251
}