from __future__ import absolute_import, division, print_function import numpy from mmtbx.tls.utils import TLSMatrices from mmtbx.tls.optimise_amplitudes import \ MultiGroupMultiDatasetUijAmplitudeFunctionalAndGradientCalculator, \ OptimiseAmplitudes, OptimisationWeights from scitbx.array_family import flex from libtbx.test_utils import approx_equal from pytest import raises from six.moves import zip from six.moves import range rran = numpy.random.random iran = numpy.random.choice numpy.random.seed(123232) TLSMatrices.set_precision(12) T_ONLY = [1.,1.,1.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.] L_ONLY = [0.,0.,0.,0.,0.,0.,1.,1.,1.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.] TLS_MATRICES = [ [0.88875254,0.266694873,0.493991604,-0.280846445,0.018132653,0.128202717,0.109792777,0.143780355,0.111757325,0.00745827,0.044408043,0.004683883,-0.031880059,0.087486387,-0.035440044,0.053723107,-0.093285664,-0.126741771,0.078707102,0.047063985,0.125165723], [0.244377462,0.459148081,0.710013128,-0.008898371,0.07279876,0.036277241,3.755718081,2.739489028,2.134579946,-0.200385262,0.074588032,-0.023382499,0.00048379,0.229469716,-0.187382678,-0.05705213,0.308356068,0.219549925,0.260721724,0.137599684,-0.308839858], [0.29790555,0.765985455,0.362144247,0.188051876,-0.07468928,0.176222813,1.157009238,1.603540537,1.146982202,0.076470473,0.076337332,-0.036514342,-0.057413883,-0.4930283,-0.141063091,0.199094828,0.180380061,0.564268854,0.271176005,0.192103145,-0.122966178], [0.550613174,0.642391951,0.512539873,0.186990637,-0.018948302,0.182865203,0.696794749,1.168473012,1.459255982,0.122241174,-0.070965716,-0.127405648,-0.126672963,-0.057326201,0.307248593,0.23402513,-0.081293525,-0.457364662,-0.270995819,0.131106455,0.207966488], [0.801839264,0.111567149,0.473172234,0.111448705,0.01322435,0.053023611,0.119600352,0.158201942,0.501511288,-0.00097325,0.038852731,-0.020179123,0.094859837,0.027409079,0.027958886,-0.036977684,0.115684103,-0.023249574,-0.167599437,0.028274946,-0.21054394], [0.476189761,0.185513316,0.494838237,0.079787304,0.186996542,0.089705044,6.468559896,0.03302114,1.360918838,-0.092774147,-0.470033825,0.142680478,0.007373318,0.322513481,-0.087204939,-0.102689013,0.085596047,-0.037679444,-0.183875771,-0.048267033,-0.092969365], [0.271946335,0.353828607,0.011539801,-0.012010653,-0.03669219,-0.006760052,23.89487928,3.564598018,0.086574587,-1.913464478,0.89104049,-0.011635907,0.00083053,-0.00120248,-0.06596398,0.059314995,0.042902408,0.001580523,-0.013344186,0.103940739,-0.043732938], [0.436887477,0.558232952,0.172244124,0.292168093,-0.130328728,-0.074615062,3.137593554,0.025205525,0.920908537,-0.108005337,0.00472811,0.015510299,-0.000822308,0.118599046,-0.151446932,0.029991796,0.042430554,-0.009368142,0.000888338,0.153578145,-0.041608246], [0.507487457,0.513473662,0.465150239,-0.024656406,0.052695005,0.030032721,9.358808433,1.113449297,12.326819619,-0.988256782,1.235061392,-1.371684512,0.454760772,-0.522429055,0.566566556,0.34528893,-0.057886461,-0.28449733,-0.171754982,-0.103504915,-0.396874311], [0.708808312,0.55187043,0.429938391,0.197257684,-3.9774e-05,-0.054427743,0.178345211,0.297789551,0.087920317,-0.003702672,-0.002424311,0.000192857,-0.029460632,0.037160795,0.061827225,0.276788373,0.000147768,-0.235898673,0.018602792,-0.125798933,0.029312864], [0.151113427,0.654574237,0.240882778,-0.018212974,-0.014025963,0.100180189,0.285310135,2.881490187,11.011691132,-0.437653779,-0.698453455,0.157185441,-0.1905967,-0.166253585,-0.151567002,0.150109019,0.444896847,-0.230918972,0.078079958,0.50626905,-0.254300147], ] def get_optimisation_test_set( n_grp, n_dst, n_atm, atomic_amplitude=1.0, random_dataset_order=False, ): """Generate a test set of random uijs for optimisation""" assert n_grp > 0 assert n_grp <= len(TLS_MATRICES) assert n_dst > 0 assert n_atm > 3 target_uijs = numpy.zeros((n_dst, n_atm, 6)) # Default to equal weights target_weights = flex.double(flex.grid((n_dst, n_atm)), 1.0) # Create atomic base atomic_values_numpy = rran(n_atm*6).reshape((n_atm,6)) real_atomic_amps = atomic_amplitude * rran(n_atm) atomic_base = flex.sym_mat3_double(atomic_values_numpy) # Create atomic total and add to target atomic_uijs = real_atomic_amps.reshape((n_atm,1)) * atomic_values_numpy for i_dst in range(n_dst): target_uijs[i_dst] = atomic_uijs # Generate a random set of amplitudes real_group_amps = rran(n_grp*n_dst).reshape((n_grp,n_dst)) # Output lists of base uijs and atoms they are associated with base_uijs = [] base_sels = [] # Mapping of elements to datasets if random_dataset_order is True: dataset_hash = flex.size_t(numpy.concatenate([flex.random_permutation(n_dst) for _ in range(n_grp)])) else: dataset_hash = flex.size_t(list(range(n_dst)) * n_grp) # base counter i_cnt = 0 for i_grp in range(n_grp): # Number of atoms in this group -- at least 2 if n_grp == 1: n_atm_this = n_atm else: n_atm_this = max(2,iran(n_atm)) # Which atoms are covered by this group i_sel = iran(n_atm, size=n_atm_this, replace=False) # Generate some uijs for this tls_m = TLSMatrices(TLS_MATRICES[i_grp]) for _ in range(n_dst): # Which dataset? i_dst = dataset_hash[i_cnt] i_cnt += 1 # Generate random coordinates x_ = 50.0*(-0.5+rran(n_atm_this*3).reshape((n_atm_this,3))) x_ = flex.vec3_double(x_) # Generate the base elements u_ = tls_m.uijs(x_, origin=(0.,0.,0.)) base_uijs.append(flex.sym_mat3_double(u_)) base_sels.append(flex.size_t(i_sel)) # Generate the amplitude-multiplied equivalent u_m = (real_group_amps[i_grp,i_dst]*tls_m).uijs(x_, origin=(0.,0.,0.)) target_uijs[i_dst, i_sel, :] += numpy.array(u_m) # Reshape target_uijs = flex.sym_mat3_double(target_uijs.reshape((n_dst*n_atm,6))) target_uijs.reshape(flex.grid((n_dst,n_atm))) return ( target_uijs, target_weights, \ base_uijs, base_sels, dataset_hash, \ atomic_base, \ real_group_amps.flatten(), real_atomic_amps, ) def resort_amplitudes_by_dataset_hash( n_grp, n_dst, dataset_hash, real_group_amps, ): """Sorted the real amplitudes into the order that they will be returned by the optimiser""" sorted_group_amps = [] i_cnt = 0 for i_grp in range(n_grp): for _ in range(n_dst): i_dst = dataset_hash[i_cnt] i_cnt += 1 sorted_group_amps.append(real_group_amps[i_grp*n_dst+i_dst]) return sorted_group_amps def calculate_expected_f_and_g_least_squares( n_grp, n_dst, n_atm, target_uijs, target_weights, base_uijs, base_sels, dataset_hash, atomic_base, current_amplitudes_base, atomic_mask = None, ): """Calculate what we expect from the current functional and gradients""" # Atomics always start with unitary amplitudes current_amplitudes_atomic = flex.double(n_atm, 1.0) if atomic_mask is None: atomic_mask = flex.bool(n_dst, True) functional = 0.0 gradients = numpy.zeros(n_grp*n_dst+n_atm) target_uijs_numpy = numpy.array(target_uijs).reshape((n_dst,n_atm,6)) target_weights_numpy = numpy.array(target_weights).reshape((n_dst,n_atm)) atomic_uijs_numpy = numpy.array(current_amplitudes_atomic).reshape((n_atm,1)) * numpy.array(atomic_base) for i_dst in range(n_dst): b_sel = [i_b for i_b, i_d in enumerate(dataset_hash) if i_d==i_dst] total_uij = numpy.zeros((n_atm, 6)) # Add up contributions from different components for i_b in b_sel: amp = current_amplitudes_base[i_b] uijs = base_uijs[i_b] atms = base_sels[i_b] assert len(uijs) == len(atms) for i_atm, u in zip(atms, uijs): total_uij[i_atm] += amp*numpy.array(u) total_uij += atomic_uijs_numpy delta_uijs = (target_uijs_numpy[i_dst] - total_uij) # Calculate functional (least squares) for u, w in zip(delta_uijs, target_weights_numpy[i_dst]): assert len(u) == 6 functional += numpy.sum(w * u * u) # Calculate gradients (for each base element) for i_b in b_sel: uijs = base_uijs[i_b] atms = base_sels[i_b] assert len(uijs) == len(atms) for i_atm, ug in zip(atms, uijs): ud = delta_uijs[i_atm] w = target_weights_numpy[i_dst, i_atm] gradients[i_b] += numpy.sum(-2.0 * w * ud * ug) if atomic_mask[i_dst]: for i_atm in range(n_atm): ur = atomic_base[i_atm] ud = delta_uijs[i_atm] w = target_weights_numpy[i_dst, i_atm] w2 = float(n_dst) / sum(atomic_mask) gradients[-n_atm+i_atm] += numpy.sum(-2.0 * w * ud * ur * w2) return functional, gradients def calculate_expected_f_and_g_sum_of_amplitudes( n_grp, n_dst, n_atm, target_weights, dataset_hash, current_amplitudes, atomic_base, overall_weight, atomic_mask = None, ): mult = overall_weight dataset_weights = target_weights.as_numpy_array().mean(axis=1) assert len(dataset_weights) == n_dst amplitude_sums = numpy.zeros(n_dst) assert len(current_amplitudes) == n_grp*n_dst + n_atm for i, a in enumerate(current_amplitudes[:n_grp*n_dst]): amplitude_sums[dataset_hash[i]] += a atomic_amps = (numpy.array(current_amplitudes[-n_atm:]).reshape((n_atm,1)) * numpy.array(atomic_base)) atomic_amps = atomic_amps[:,0:3].mean(axis=1) assert len(atomic_amps) == n_atm # Add to all datasets amplitude_sums[:] += atomic_amps.sum() f = mult * (amplitude_sums * dataset_weights).sum() g = flex.double( [mult * dataset_weights[i_dst] for i_dst in dataset_hash] + [mult * dataset_weights.sum()]*n_atm ) return f,g def calculate_expected_f_and_g_sum_of_amplitudes_squared( n_grp, n_dst, n_atm, target_weights, dataset_hash, current_amplitudes, atomic_base, overall_weight, atomic_mask = None, ): mult = overall_weight dataset_weights = target_weights.as_numpy_array().mean(axis=1) assert len(dataset_weights) == n_dst amplitude_sums = numpy.zeros(n_dst) assert len(current_amplitudes) == n_grp*n_dst + n_atm for i, a in enumerate(current_amplitudes[:n_grp*n_dst]): amplitude_sums[dataset_hash[i]] += a atomic_amps = (numpy.array(current_amplitudes[-n_atm:]).reshape((n_atm,1)) * numpy.array(atomic_base)) atomic_amps = atomic_amps[:,0:3].mean(axis=1) assert len(atomic_amps) == n_atm # Add to all datasets amplitude_sums[:] += atomic_amps.sum() f = mult * (amplitude_sums * amplitude_sums * dataset_weights).sum() g = flex.double( [2.0 * mult * amplitude_sums[i_dst] * dataset_weights[i_dst] for i_dst in dataset_hash] + [2.0 * mult * (amplitude_sums * dataset_weights).sum()]*n_atm ) return f,g def calculate_expected_f_and_g_sum_of_squared_amplitudes( n_grp, n_dst, n_atm, target_weights, dataset_hash, current_amplitudes, atomic_base, overall_weight, atomic_mask=None, ): mult = overall_weight dataset_weights = target_weights.as_numpy_array().mean(axis=1) assert len(dataset_weights) == n_dst dataset_weights_base = numpy.array([dataset_weights[i_dst] for i_dst in dataset_hash]) assert len(current_amplitudes) == n_grp*n_dst + n_atm group_amplitudes = current_amplitudes[:n_grp*n_dst] atomic_amps = (numpy.array(current_amplitudes[-n_atm:]).reshape((n_atm,1)) * numpy.array(atomic_base)) atomic_amps = atomic_amps[:,0:3].mean(axis=1) assert len(atomic_amps) == n_atm f = mult * ( (group_amplitudes * group_amplitudes * dataset_weights_base).sum() + (atomic_amps * atomic_amps * dataset_weights.sum()).sum() ) g = flex.double( list(2.0 * mult * group_amplitudes * dataset_weights_base) + list(2.0 * mult * atomic_amps * dataset_weights.sum()) ) return f,g def tst_optimisation_weights(): wgts_dbl = rran(3) wgts_dict = dict( sum_of_amplitudes = wgts_dbl[0], sum_of_squared_amplitudes = wgts_dbl[1], sum_of_amplitudes_squared = wgts_dbl[2], ) # Initialise from dict weights = OptimisationWeights(**wgts_dict) assert approx_equal(weights.sum_of_amplitudes, wgts_dbl[0]) assert approx_equal(weights.sum_of_squared_amplitudes, wgts_dbl[1]) assert approx_equal(weights.sum_of_amplitudes_squared, wgts_dbl[2]) print('OK') def tst_optimise_amplitudes_single_group(): """Check that very basic optimisation reaches expected values regardless of starting values""" n_grp = 1 n_atm = 10 for n_dst in [1,5]: for g_amp_start in [0.0, 1.0, None]: for r_amp_mult in [0.0, 1.0]: # Get a random set of values to test target_uijs, target_weights, \ base_uijs, base_sels, dataset_hash, \ atomic_base, \ real_group_amps, real_atomic_amps \ = get_optimisation_test_set(n_grp, n_dst, n_atm, atomic_amplitude=r_amp_mult) if r_amp_mult == 0.0: assert not real_atomic_amps.any() # Starting values if g_amp_start is not None: base_amplitudes_start = flex.double(n_grp*n_dst, g_amp_start) else: base_amplitudes_start = flex.double(rran(n_grp*n_dst)) # Since not given, amplitudes start at 1.0 for atomic atomic_amplitudes_start = flex.double(n_atm, 1.0) # Optimise opt = OptimiseAmplitudes( target_uijs = target_uijs, target_weights = target_weights, base_amplitudes = base_amplitudes_start, base_uijs = base_uijs, base_atom_indices = base_sels, base_dataset_hash = dataset_hash, atomic_uijs = atomic_base, atomic_amplitudes = None, atomic_optimisation_mask = None, optimisation_weights = None, convergence_tolerance=1e-08, ).run() assert approx_equal(list(opt.initial), list(base_amplitudes_start) + list(atomic_amplitudes_start), 1e-6) assert approx_equal(list(opt.result), list(real_group_amps) + list(real_atomic_amps), 1e-6) print('OK') def tst_optimise_amplitudes_multiple_groups_with_atomic(): """Check that multiple partial groups optimise correctly with atomic level""" n_grp = 3 n_dst = 5 n_atm = 10 # Get a random set of values to test target_uijs, target_weights, \ base_uijs, base_sels, dataset_hash, \ atomic_base, \ real_group_amps, real_atomic_amps \ = get_optimisation_test_set(n_grp, n_dst, n_atm) # Starting values base_amplitudes_start = flex.double(n_grp*n_dst, 0.0) atomic_amplitudes_start = flex.double(n_atm, 0.0) opt = OptimiseAmplitudes( target_uijs = target_uijs, target_weights = target_weights, base_amplitudes = base_amplitudes_start, base_uijs = base_uijs, base_atom_indices = base_sels, base_dataset_hash = dataset_hash, atomic_uijs = atomic_base, atomic_amplitudes = atomic_amplitudes_start, atomic_optimisation_mask = None, optimisation_weights = None, convergence_tolerance=1e-08, ).run() assert approx_equal(list(opt.initial), list(base_amplitudes_start) + list(atomic_amplitudes_start), 1e-6) assert approx_equal(list(opt.result), list(real_group_amps) + list(real_atomic_amps), 1e-6) ######################################################## # Check that shuffling the input lists has the expected effect ######################################################## n_base = n_grp * n_dst i_perm = iran(n_base, size=n_base, replace=False) # Reorder the base elements by random permutation base_uijs = [base_uijs[i] for i in i_perm] base_sels = [base_sels[i] for i in i_perm] dataset_hash = flex.size_t([dataset_hash[i] for i in i_perm]) base_amplitudes_start = flex.double([base_amplitudes_start[i] for i in i_perm]) opt = OptimiseAmplitudes( target_uijs = target_uijs, target_weights = target_weights, base_amplitudes = base_amplitudes_start, base_uijs = base_uijs, base_atom_indices = base_sels, base_dataset_hash = dataset_hash, atomic_uijs = atomic_base, atomic_amplitudes = atomic_amplitudes_start, atomic_optimisation_mask = None, optimisation_weights = None, convergence_tolerance=1e-08, ).run() assert approx_equal(list(opt.initial), list(base_amplitudes_start) + list(atomic_amplitudes_start), 1e-6) assert approx_equal(list(opt.result), list([real_group_amps[i] for i in i_perm]) + list(real_atomic_amps), 1e-6) print('OK') def tst_optimise_amplitudes_multiple_groups_permuted_dataset_order(): n_grp = 3 n_dst = 5 n_atm = 10 # Get a random set of values to test target_uijs, target_weights, \ base_uijs, base_sels, dataset_hash, \ atomic_base, \ real_group_amps, real_atomic_amps \ = get_optimisation_test_set(n_grp, n_dst, n_atm, random_dataset_order=True) sorted_group_amps = resort_amplitudes_by_dataset_hash( n_grp=n_grp, n_dst=n_dst, dataset_hash=dataset_hash, real_group_amps=real_group_amps, ) # Starting values base_amplitudes_start = flex.double(n_grp*n_dst, 0.0) atomic_amplitudes_start = flex.double(n_atm, 0.0) opt = OptimiseAmplitudes( target_uijs = target_uijs, target_weights = target_weights, base_amplitudes = base_amplitudes_start, base_uijs = base_uijs, base_atom_indices = base_sels, base_dataset_hash = dataset_hash, atomic_uijs = atomic_base, atomic_amplitudes = atomic_amplitudes_start, atomic_optimisation_mask = None, optimisation_weights = None, convergence_tolerance=1e-08, ).run() assert approx_equal(list(opt.initial), list(base_amplitudes_start) + list(atomic_amplitudes_start), 1e-6) assert approx_equal(list(opt.result), list(sorted_group_amps) + list(real_atomic_amps), 1e-6) print('OK') def tst_functional_gradient_calculator(): """Check that the functional and gradients are calculated as expected""" n_grp = 3 n_dst = 5 n_atm = 10 for atomic_mask_size in [None, 1, 2, 3]: if atomic_mask_size is not None: atomic_mask = flex.bool(n_dst, False) n_sel = max(n_dst, atomic_mask_size) atomic_mask.set_selected(flex.size_t(iran(n_dst, size=n_sel, replace=False)), True) assert sum(atomic_mask) == n_sel else: atomic_mask = None target_uijs, target_weights, \ base_uijs, base_sels, dataset_hash, \ atomic_base, \ real_group_amps, real_atomic_amps \ = get_optimisation_test_set(n_grp, n_dst, n_atm) # Starting values - randomise base_amplitudes_start = flex.random_double(n_grp*n_dst) ######################################################## # Check least-squares f/g ######################################################## functional, gradients = calculate_expected_f_and_g_least_squares( n_grp = n_grp, n_dst = n_dst, n_atm = n_atm, target_uijs = target_uijs, target_weights = target_weights, base_uijs = base_uijs, base_sels = base_sels, dataset_hash = dataset_hash, atomic_base = atomic_base, current_amplitudes_base = base_amplitudes_start, atomic_mask = atomic_mask, ) f_g_calculator = MultiGroupMultiDatasetUijAmplitudeFunctionalAndGradientCalculator( target_uijs = target_uijs, target_weights = target_weights, base_amplitudes = base_amplitudes_start, base_uijs = base_uijs, base_atom_indices = base_sels, base_dataset_hash = dataset_hash, atomic_uijs = atomic_base, weight_sum_of_amplitudes = 0.0, weight_sum_of_amplitudes_squared = 0.0, weight_sum_of_squared_amplitudes = 0.0, ) if atomic_mask: f_g_calculator.set_atomic_optimisation_mask(atomic_mask) f, g = f_g_calculator.compute_functional_and_gradients() assert approx_equal(f, functional) assert approx_equal(list(g), list(gradients)) ######################################################## # check that correct solution gives no functional and gradient ######################################################## correct_amplitudes = flex.double(list(real_group_amps) + list(real_atomic_amps)) # Manually set amplitudes f_g_calculator.set_current_amplitudes(correct_amplitudes) f, g = f_g_calculator.compute_functional_and_gradients() assert approx_equal(f, 0.0) assert g.all_approx_equal(0.0) ######################################################## # Calculate gradients for weight_sum_of_amplitudes ######################################################## # !!! Still using the correct solution so that those functional and gradient components are zero wgt = rran() f_g_calculator = MultiGroupMultiDatasetUijAmplitudeFunctionalAndGradientCalculator( target_uijs = target_uijs, target_weights = target_weights, base_amplitudes = base_amplitudes_start, base_uijs = base_uijs, base_atom_indices = base_sels, base_dataset_hash = dataset_hash, atomic_uijs = atomic_base, weight_sum_of_amplitudes = wgt, weight_sum_of_amplitudes_squared = 0.0, weight_sum_of_squared_amplitudes = 0.0, ) if atomic_mask: f_g_calculator.set_atomic_optimisation_mask(atomic_mask) f_g_calculator.set_current_amplitudes(correct_amplitudes) functional, gradients = calculate_expected_f_and_g_sum_of_amplitudes( n_grp = n_grp, n_dst = n_dst, n_atm = n_atm, target_weights = target_weights, dataset_hash = dataset_hash, current_amplitudes = f_g_calculator.get_current_amplitudes(), atomic_base = atomic_base, overall_weight = wgt, ) f, g = f_g_calculator.compute_functional_and_gradients() assert approx_equal(f, functional) assert approx_equal(list(g), list(gradients)) ######################################################## # Calculate gradients for weight_sum_of_amplitudes_squared ######################################################## # !!! Still using the correct solution so that those functional and gradient components are zero wgt = rran() f_g_calculator = MultiGroupMultiDatasetUijAmplitudeFunctionalAndGradientCalculator( target_uijs = target_uijs, target_weights = target_weights, base_amplitudes = base_amplitudes_start, base_uijs = base_uijs, base_atom_indices = base_sels, base_dataset_hash = dataset_hash, atomic_uijs = atomic_base, weight_sum_of_amplitudes = 0.0, weight_sum_of_amplitudes_squared = wgt, weight_sum_of_squared_amplitudes = 0.0, ) if atomic_mask: f_g_calculator.set_atomic_optimisation_mask(atomic_mask) f_g_calculator.set_current_amplitudes(correct_amplitudes) functional, gradients = calculate_expected_f_and_g_sum_of_amplitudes_squared( n_grp = n_grp, n_dst = n_dst, n_atm = n_atm, target_weights = target_weights, dataset_hash = dataset_hash, current_amplitudes = f_g_calculator.get_current_amplitudes(), atomic_base = atomic_base, overall_weight = wgt, ) f, g = f_g_calculator.compute_functional_and_gradients() assert approx_equal(f, functional) assert approx_equal(list(g), list(gradients)) ######################################################## # Calculate gradients for weight_sum_of_squared_amplitudes ######################################################## # !!! Still using the correct solution so that those functional and gradient components are zero wgt = rran() f_g_calculator = MultiGroupMultiDatasetUijAmplitudeFunctionalAndGradientCalculator( target_uijs = target_uijs, target_weights = target_weights, base_amplitudes = base_amplitudes_start, base_uijs = base_uijs, base_atom_indices = base_sels, base_dataset_hash = dataset_hash, atomic_uijs = atomic_base, weight_sum_of_amplitudes = 0.0, weight_sum_of_amplitudes_squared = 0.0, weight_sum_of_squared_amplitudes = wgt, ) if atomic_mask: f_g_calculator.set_atomic_optimisation_mask(atomic_mask) f_g_calculator.set_current_amplitudes(correct_amplitudes) functional, gradients = calculate_expected_f_and_g_sum_of_squared_amplitudes( n_grp = n_grp, n_dst = n_dst, n_atm = n_atm, target_weights = target_weights, dataset_hash = dataset_hash, current_amplitudes = f_g_calculator.get_current_amplitudes(), atomic_base = atomic_base, overall_weight = wgt, ) f, g = f_g_calculator.compute_functional_and_gradients() assert approx_equal(f, functional) assert approx_equal(list(g), list(gradients)) ######################################################## # Turn on all weights and non-optimal solution ######################################################## wgts = rran(3) f_g_calculator = MultiGroupMultiDatasetUijAmplitudeFunctionalAndGradientCalculator( target_uijs = target_uijs, target_weights = target_weights, base_amplitudes = base_amplitudes_start, base_uijs = base_uijs, base_atom_indices = base_sels, base_dataset_hash = dataset_hash, atomic_uijs = atomic_base, weight_sum_of_amplitudes = wgts[0], weight_sum_of_amplitudes_squared = wgts[1], weight_sum_of_squared_amplitudes = wgts[2], ) if atomic_mask: f_g_calculator.set_atomic_optimisation_mask(atomic_mask) flsq, glsq = calculate_expected_f_and_g_least_squares( n_grp = n_grp, n_dst = n_dst, n_atm = n_atm, target_uijs = target_uijs, target_weights = target_weights, base_uijs = base_uijs, base_sels = base_sels, dataset_hash = dataset_hash, atomic_base = atomic_base, current_amplitudes_base = base_amplitudes_start, atomic_mask = atomic_mask, ) fw1, gw1 = calculate_expected_f_and_g_sum_of_amplitudes( n_grp = n_grp, n_dst = n_dst, n_atm = n_atm, target_weights = target_weights, dataset_hash = dataset_hash, current_amplitudes = f_g_calculator.get_current_amplitudes(), atomic_base = atomic_base, overall_weight = wgts[0], ) fw2, gw2 = calculate_expected_f_and_g_sum_of_amplitudes_squared( n_grp = n_grp, n_dst = n_dst, n_atm = n_atm, target_weights = target_weights, dataset_hash = dataset_hash, current_amplitudes = f_g_calculator.get_current_amplitudes(), atomic_base = atomic_base, overall_weight = wgts[1], ) fw3, gw3 = calculate_expected_f_and_g_sum_of_squared_amplitudes( n_grp = n_grp, n_dst = n_dst, n_atm = n_atm, target_weights = target_weights, dataset_hash = dataset_hash, current_amplitudes = f_g_calculator.get_current_amplitudes(), atomic_base = atomic_base, overall_weight = wgts[2], ) f, g = f_g_calculator.compute_functional_and_gradients() assert approx_equal(f, flsq+fw1+fw2+fw3) assert approx_equal(list(g), list(glsq+gw1+gw2+gw3)) ######################################################## # Make some amplitudes negative and recalculate - negative amplitudes should be set to zero ######################################################## # Starting values base_amplitudes_start = flex.double(n_grp*n_dst, 0.0) functional, gradients = calculate_expected_f_and_g_least_squares( n_grp = n_grp, n_dst = n_dst, n_atm = n_atm, target_uijs = target_uijs, target_weights = target_weights, base_uijs = base_uijs, base_sels = base_sels, dataset_hash = dataset_hash, atomic_base = atomic_base, current_amplitudes_base = base_amplitudes_start, atomic_mask = atomic_mask, ) f_g_calculator = MultiGroupMultiDatasetUijAmplitudeFunctionalAndGradientCalculator( target_uijs = target_uijs, target_weights = target_weights, base_amplitudes = base_amplitudes_start, base_uijs = base_uijs, base_atom_indices = base_sels, base_dataset_hash = dataset_hash, atomic_uijs = atomic_base, weight_sum_of_amplitudes = 0.0, weight_sum_of_amplitudes_squared = 0.0, weight_sum_of_squared_amplitudes = 0.0, ) if atomic_mask: f_g_calculator.set_atomic_optimisation_mask(atomic_mask) f, g = f_g_calculator.compute_functional_and_gradients() assert approx_equal(f, functional) assert approx_equal(list(g), list(gradients)) # Now make some amplitudes negative base_amplitudes_start.set_selected(flex.size_t(iran(n_grp*n_dst)), -0.5) base_amplitudes_start.set_selected(flex.size_t(iran(n_grp*n_dst)), -200.) f_g_calculator = MultiGroupMultiDatasetUijAmplitudeFunctionalAndGradientCalculator( target_uijs = target_uijs, target_weights = target_weights, base_amplitudes = base_amplitudes_start, base_uijs = base_uijs, base_atom_indices = base_sels, base_dataset_hash = dataset_hash, atomic_uijs = atomic_base, weight_sum_of_amplitudes = 0.0, weight_sum_of_amplitudes_squared = 0.0, weight_sum_of_squared_amplitudes = 0.0, ) if atomic_mask: f_g_calculator.set_atomic_optimisation_mask(atomic_mask) f, g = f_g_calculator.compute_functional_and_gradients() assert approx_equal(f, functional) assert approx_equal(list(g), list(gradients)) print('OK') def tst_functional_gradient_calculator_permuted_dataset_order(): n_grp = 3 n_dst = 5 n_atm = 10 for atomic_mask_size in [None, 1, 2, 3]: if atomic_mask_size is not None: atomic_mask = flex.bool(n_dst, False) n_sel = max(n_dst, atomic_mask_size) atomic_mask.set_selected(flex.size_t(iran(n_dst, size=n_sel, replace=False)), True) assert sum(atomic_mask) == n_sel else: atomic_mask = None target_uijs, target_weights, \ base_uijs, base_sels, dataset_hash, \ atomic_base, \ real_group_amps, real_atomic_amps \ = get_optimisation_test_set(n_grp, n_dst, n_atm, random_dataset_order=True) # Starting values - randomise base_amplitudes_start = flex.random_double(n_grp*n_dst) wgts = rran(3) f_g_calculator = MultiGroupMultiDatasetUijAmplitudeFunctionalAndGradientCalculator( target_uijs = target_uijs, target_weights = target_weights, base_amplitudes = base_amplitudes_start, base_uijs = base_uijs, base_atom_indices = base_sels, base_dataset_hash = dataset_hash, atomic_uijs = atomic_base, weight_sum_of_amplitudes = wgts[0], weight_sum_of_amplitudes_squared = wgts[1], weight_sum_of_squared_amplitudes = wgts[2], ) if atomic_mask: f_g_calculator.set_atomic_optimisation_mask(atomic_mask) flsq, glsq = calculate_expected_f_and_g_least_squares( n_grp = n_grp, n_dst = n_dst, n_atm = n_atm, target_uijs = target_uijs, target_weights = target_weights, base_uijs = base_uijs, base_sels = base_sels, dataset_hash = dataset_hash, atomic_base = atomic_base, current_amplitudes_base = base_amplitudes_start, atomic_mask = atomic_mask, ) fw1, gw1 = calculate_expected_f_and_g_sum_of_amplitudes( n_grp = n_grp, n_dst = n_dst, n_atm = n_atm, target_weights = target_weights, dataset_hash = dataset_hash, current_amplitudes = f_g_calculator.get_current_amplitudes(), atomic_base = atomic_base, overall_weight = wgts[0], ) fw2, gw2 = calculate_expected_f_and_g_sum_of_amplitudes_squared( n_grp = n_grp, n_dst = n_dst, n_atm = n_atm, target_weights = target_weights, dataset_hash = dataset_hash, current_amplitudes = f_g_calculator.get_current_amplitudes(), atomic_base = atomic_base, overall_weight = wgts[1], ) fw3, gw3 = calculate_expected_f_and_g_sum_of_squared_amplitudes( n_grp = n_grp, n_dst = n_dst, n_atm = n_atm, target_weights = target_weights, dataset_hash = dataset_hash, current_amplitudes = f_g_calculator.get_current_amplitudes(), atomic_base = atomic_base, overall_weight = wgts[2], ) f, g = f_g_calculator.compute_functional_and_gradients() assert approx_equal(f, flsq+fw1+fw2+fw3) assert approx_equal(list(g), list(glsq+gw1+gw2+gw3)) print('OK') def tst_functional_gradient_calculator_invalid_arguments(): """Check errors are raised as expected""" n_grp = 3 n_dst = 5 n_atm = 10 target_uijs, target_weights, \ base_uijs, base_sels, dataset_hash, \ atomic_base, \ real_group_amps, real_atomic_amps \ = get_optimisation_test_set(n_grp, n_dst, n_atm) ######################################################## # Check expected error messages are raised ######################################################## # Starting values base_amplitudes_start = flex.double(n_grp*n_dst, 1.0) wgt_kw_args = dict( weight_sum_of_amplitudes = 0.0, weight_sum_of_amplitudes_squared = 0.0, weight_sum_of_squared_amplitudes = 0.0, ) # Should not error f_g_calculator = MultiGroupMultiDatasetUijAmplitudeFunctionalAndGradientCalculator( target_uijs = target_uijs, target_weights = target_weights, base_amplitudes = base_amplitudes_start, base_uijs = base_uijs, base_atom_indices = base_sels, base_dataset_hash = dataset_hash, atomic_uijs = atomic_base, **wgt_kw_args ) f, g = f_g_calculator.compute_functional_and_gradients() # target_uijs msg = "invalid target_uijs: must be 2-dimensional flex array (currently 3)" with raises(Exception) as e: f_g_calculator = MultiGroupMultiDatasetUijAmplitudeFunctionalAndGradientCalculator( target_uijs = flex.sym_mat3_double(flex.grid((n_dst-1,n_atm,5)), (1.,1.,1.,0.,0.,0.)), target_weights = target_weights, base_amplitudes = base_amplitudes_start, base_uijs = base_uijs, base_atom_indices = base_sels, base_dataset_hash = dataset_hash, atomic_uijs = atomic_base, **wgt_kw_args ) assert msg == str(e.value), '"{}" does not match "{}"'.format(msg, str(e.value)) # target_weights msg = "invalid dimension of target_weights (dimension 3): must be same dimension as target_uijs (dimension 2)" with raises(Exception) as e: f_g_calculator = MultiGroupMultiDatasetUijAmplitudeFunctionalAndGradientCalculator( target_uijs = target_uijs, target_weights = flex.double(flex.grid((n_dst,n_atm,5)), 1.0), base_amplitudes = base_amplitudes_start, base_uijs = base_uijs, base_atom_indices = base_sels, base_dataset_hash = dataset_hash, atomic_uijs = atomic_base, **wgt_kw_args ) assert msg == str(e.value) msg = "incompatible dimension of target_weights (axis 0): must be same size as target_uijs ({} != {})".format(n_dst, n_dst-1) with raises(Exception) as e: f_g_calculator = MultiGroupMultiDatasetUijAmplitudeFunctionalAndGradientCalculator( target_uijs = flex.sym_mat3_double(flex.grid((n_dst-1,n_atm)), (1.,1.,1.,0.,0.,0.)), target_weights = target_weights, base_amplitudes = base_amplitudes_start, base_uijs = base_uijs, base_atom_indices = base_sels, base_dataset_hash = dataset_hash, atomic_uijs = atomic_base, **wgt_kw_args ) assert msg == str(e.value) msg = "incompatible dimension of target_weights (axis 1): must be same size as target_uijs ({} != {})".format(n_atm, n_atm+1) with raises(Exception) as e: f_g_calculator = MultiGroupMultiDatasetUijAmplitudeFunctionalAndGradientCalculator( target_uijs = flex.sym_mat3_double(flex.grid((n_dst,n_atm+1)), (1.,1.,1.,0.,0.,0.)), target_weights = target_weights, base_amplitudes = base_amplitudes_start, base_uijs = base_uijs, base_atom_indices = base_sels, base_dataset_hash = dataset_hash, atomic_uijs = atomic_base, **wgt_kw_args ) assert msg == str(e.value) # base components msg = "invalid input base components. base_amplitudes (length {}), base_uijs (length {}) and base_atom_indices (length {}) must all be the same length" with raises(Exception) as e: f_g_calculator = MultiGroupMultiDatasetUijAmplitudeFunctionalAndGradientCalculator( target_uijs = target_uijs, target_weights = target_weights, base_amplitudes = flex.double(n_grp*n_dst-1, 1.0), base_uijs = base_uijs, base_atom_indices = base_sels, base_dataset_hash = dataset_hash, atomic_uijs = atomic_base, **wgt_kw_args ) assert msg.format(n_grp*n_dst-1, n_grp*n_dst, n_grp*n_dst) == str(e.value) with raises(Exception) as e: f_g_calculator = MultiGroupMultiDatasetUijAmplitudeFunctionalAndGradientCalculator( target_uijs = target_uijs, target_weights = target_weights, base_amplitudes = base_amplitudes_start, base_uijs = base_uijs[:-1], base_atom_indices = base_sels[:-1], base_dataset_hash = dataset_hash, atomic_uijs = atomic_base, **wgt_kw_args ) assert msg.format(n_grp*n_dst, n_grp*n_dst-1, n_grp*n_dst-1) == str(e.value) with raises(Exception) as e: f_g_calculator = MultiGroupMultiDatasetUijAmplitudeFunctionalAndGradientCalculator( target_uijs = target_uijs, target_weights = target_weights, base_amplitudes = base_amplitudes_start, base_uijs = base_uijs[:-1], base_atom_indices = base_sels, base_dataset_hash = dataset_hash, atomic_uijs = atomic_base, **wgt_kw_args ) assert msg.format(n_grp*n_dst, n_grp*n_dst-1, n_grp*n_dst) == str(e.value) with raises(Exception) as e: f_g_calculator = MultiGroupMultiDatasetUijAmplitudeFunctionalAndGradientCalculator( target_uijs = target_uijs, target_weights = target_weights, base_amplitudes = base_amplitudes_start, base_uijs = base_uijs, base_atom_indices = base_sels[:-1], base_dataset_hash = dataset_hash, atomic_uijs = atomic_base, **wgt_kw_args ) assert msg.format(n_grp*n_dst, n_grp*n_dst, n_grp*n_dst-1) == str(e.value) msg = "incompatible pair (element 2) in base_uijs/base_atom_indices: pairwise elements must be the same length ({} and {})".format(n_atm+1, len(base_uijs[2])) with raises(Exception) as e: f_g_calculator = MultiGroupMultiDatasetUijAmplitudeFunctionalAndGradientCalculator( target_uijs = target_uijs, target_weights = target_weights, base_amplitudes = base_amplitudes_start, base_uijs = base_uijs[:2] + [flex.sym_mat3_double(n_atm+1)] + base_uijs[3:], base_atom_indices = base_sels, base_dataset_hash = dataset_hash, atomic_uijs = atomic_base, **wgt_kw_args ) assert msg == str(e.value) msg = "incompatible pair (element 2) in base_uijs/base_atom_indices: pairwise elements must be the same length ({} and {})".format(len(base_sels[2]), n_atm+1) with raises(Exception) as e: f_g_calculator = MultiGroupMultiDatasetUijAmplitudeFunctionalAndGradientCalculator( target_uijs = target_uijs, target_weights = target_weights, base_amplitudes = base_amplitudes_start, base_uijs = base_uijs, base_atom_indices = base_sels[:2] + [flex.size_t_range(n_atm+1)] + base_sels[3:], base_dataset_hash = dataset_hash, atomic_uijs = atomic_base, **wgt_kw_args ) assert msg == str(e.value) msg = "invalid selection in base_atom_indices ({}): attempting to select atom outside of array (size {})".format(n_atm, n_atm) with raises(Exception) as e: f_g_calculator = MultiGroupMultiDatasetUijAmplitudeFunctionalAndGradientCalculator( target_uijs = target_uijs, target_weights = target_weights, base_amplitudes = base_amplitudes_start, base_uijs = base_uijs[:2] + [flex.sym_mat3_double(n_atm+1)] + base_uijs[3:], base_atom_indices = base_sels[:2] + [flex.size_t_range(n_atm+1)] + base_sels[3:], base_dataset_hash = dataset_hash, atomic_uijs = atomic_base, **wgt_kw_args ) assert msg == str(e.value) # dataset_hash msg = "invalid base_dataset_hash (length {}): must be same length as base_amplitudes, base_uijs & base_atom_indices (length {})".format(len(dataset_hash)-1, len(dataset_hash)) with raises(Exception) as e: f_g_calculator = MultiGroupMultiDatasetUijAmplitudeFunctionalAndGradientCalculator( target_uijs = target_uijs, target_weights = target_weights, base_amplitudes = base_amplitudes_start, base_uijs = base_uijs, base_atom_indices = base_sels, base_dataset_hash = flex.size_t(list(dataset_hash)[:-1]), atomic_uijs = atomic_base, **wgt_kw_args ) assert msg == str(e.value), (msg, str(e.value)) msg = "invalid value in base_dataset_hash ({}): attempts to select element outside range of target_uijs (size {})".format(n_dst-1, n_dst-1) with raises(Exception) as e: f_g_calculator = MultiGroupMultiDatasetUijAmplitudeFunctionalAndGradientCalculator( target_uijs = flex.sym_mat3_double(flex.grid((n_dst-1,n_atm)), (1.,1.,1.,0.,0.,0.)), target_weights = flex.double(flex.grid((n_dst-1,n_atm)), 1.0), # need to resize this also base_amplitudes = base_amplitudes_start, base_uijs = base_uijs, base_atom_indices = base_sels, base_dataset_hash = dataset_hash, atomic_uijs = atomic_base, **wgt_kw_args ) assert msg == str(e.value), (msg, str(e.value)) msg = "Dataset index {} is not present in base_dataset_hash -- this dataset has no base elements associated with it.".format(n_dst) with raises(Exception) as e: f_g_calculator = MultiGroupMultiDatasetUijAmplitudeFunctionalAndGradientCalculator( target_uijs = flex.sym_mat3_double(flex.grid((n_dst+1,n_atm)), (1.,1.,1.,0.,0.,0.)), target_weights = flex.double(flex.grid((n_dst+1,n_atm)), 1.0), # need to resize this also base_amplitudes = base_amplitudes_start, base_uijs = base_uijs, base_atom_indices = base_sels, base_dataset_hash = dataset_hash, atomic_uijs = atomic_base, **wgt_kw_args ) assert msg == str(e.value), (msg, str(e.value)) # atomic uijs msg = "invalid size of atomic_uijs ({}): must match 2nd dimension of target_uijs ({})".format(n_atm-1, n_atm) with raises(Exception) as e: f_g_calculator = MultiGroupMultiDatasetUijAmplitudeFunctionalAndGradientCalculator( target_uijs = target_uijs, target_weights = target_weights, base_amplitudes = base_amplitudes_start, base_uijs = base_uijs, base_atom_indices = base_sels, base_dataset_hash = dataset_hash, atomic_uijs = flex.sym_mat3_double(n_atm-1), **wgt_kw_args ) assert msg == str(e.value), (msg, str(e.value)) # atomic mask f_g_calculator = MultiGroupMultiDatasetUijAmplitudeFunctionalAndGradientCalculator( target_uijs = target_uijs, target_weights = target_weights, base_amplitudes = base_amplitudes_start, base_uijs = base_uijs, base_atom_indices = base_sels, base_dataset_hash = dataset_hash, atomic_uijs = atomic_base, **wgt_kw_args ) # should not error f_g_calculator.set_atomic_optimisation_mask(flex.bool(n_dst, True)) f_g_calculator.set_atomic_optimisation_mask(flex.bool(n_dst, False)) # should error msg = "Input array (size {}) must be the same length as number of datasets ({})".format(n_dst-1, n_dst) with raises(Exception) as e: f_g_calculator.set_atomic_optimisation_mask(flex.bool(n_dst-1, True)) assert msg == str(e.value) msg = "Input array (size {}) must be the same length as number of datasets ({})".format(n_dst+1, n_dst) with raises(Exception) as e: f_g_calculator.set_atomic_optimisation_mask(flex.bool(n_dst+1, True)) assert msg == str(e.value) # setting amplitudes # should not error f_g_calculator.set_current_amplitudes(flex.double(n_grp*n_dst+n_atm)) # should error msg = "Input array (size {}) must be the same length as current_amplitudes (size {})".format(n_grp*n_dst+n_atm-1, n_grp*n_dst+n_atm) with raises(Exception) as e: f_g_calculator.set_current_amplitudes(flex.double(n_grp*n_dst+n_atm-1)) assert msg == str(e.value) print('OK') if __name__ == "__main__": tst_optimisation_weights() tst_optimise_amplitudes_single_group() tst_optimise_amplitudes_multiple_groups_with_atomic() tst_optimise_amplitudes_multiple_groups_permuted_dataset_order() tst_functional_gradient_calculator() tst_functional_gradient_calculator_permuted_dataset_order() tst_functional_gradient_calculator_invalid_arguments()