def write_ccp4_map(realmap, mapfile, cell, spacegroup):
"""
Write CCP4 map file from NumPy array of real-space density.
Parameters
----------
realmap : np.ndarray
3D NumPy array of real-space, voxelized electron density
mapfile : str
Filename to which map will be written
cell : gemmi.UnitCell
Unit cell parameters to use in map file
spacegroup : gemmi.SpaceGroup
Spacegroup to use in map file
"""
if not isinstance(realmap, np.ndarray) or not (realmap.ndim == 3):
raise ValueError("realmap must be a 3-dimension NumPy array")
# Set up gemmi FloatGrid object with NumPy array
grid = gemmi.FloatGrid(*realmap.shape)
grid.set_unit_cell(cell)
grid.spacegroup = spacegroup
temp = np.array(grid, copy=False)
temp[:, :, :] = realmap[:, :, :]
# Write CCP4 map
ccp4 = gemmi.Ccp4Map()
ccp4.grid = grid
ccp4.update_ccp4_header(2, True)
ccp4.write_ccp4_map(mapfile)
return
def make_mean_map(grids: List[gemmi.FloatGrid]):
reference_xmap = grids[0]
arrays_list = [np.array(grid, copy=False) for grid in grids]
arrays = np.stack(arrays_list, axis=0)
mean_array = np.mean(arrays, axis=0)
mean_grid = gemmi.FloatGrid(
reference_xmap.nu,
reference_xmap.nv,
reference_xmap.nw,
)
mean_grid.set_unit_cell(reference_xmap.unit_cell)
mean_grid.spacegroup = reference_xmap.spacegroup
for point in mean_grid:
u = point.u
v = point.v
w = point.w
value = mean_array[u, v, w]
mean_grid.set_value(u, v, w, value)
return mean_grid
def get_rscc(_reference_grid, _grid, reference_structure):
mask = gemmi.FloatGrid(_reference_grid.nu,
_reference_grid.nv,
_reference_grid.nw, )
mask.set_unit_cell(_reference_grid.unit_cell)
mask.spacegroup = gemmi.find_spacegroup_by_name("P 1")
for model in reference_structure:
for chain in model:
for residue in chain.get_polymer():
for atom in residue:
mask.set_points_around(atom.pos, 3.0, 1.0)
_reference_sample = np.array(_reference_grid, copy=False)
_sample = np.array(_grid, copy=False)
_mask = np.array(mask, copy=False)
reference_flattened = _reference_sample[_mask!= 0]
sample_flattened = _sample[_mask!=0]
reference_sample_mean = np.mean(reference_flattened)
reference_sample_demeaned = reference_flattened - reference_sample_mean
reference_sample_denominator = np.sqrt(np.sum(np.square(reference_sample_demeaned)))
sample_mean = np.mean(sample_flattened)
sample_demeaned = sample_flattened - sample_mean
sample_denominator = np.sqrt(np.sum(np.square(sample_demeaned)))
nominator = np.sum(reference_sample_demeaned * sample_demeaned)
denominator = sample_denominator * reference_sample_denominator
correlation = nominator / denominator
return correlation
def to_gemmi(self):
grid = gemmi.FloatGrid(*self.array.shape)
grid.spacegroup = self.spacegroup.to_gemmi()
grid.set_unit_cell(self.unit_cell.to_gemmi())
grid_array = np.array(grid, copy=False)
grid_array[:, :, :] = self.array[:, :, :]
return grid
def test_reading(self):
path = os.path.join(os.path.dirname(__file__), '5i55_tiny.ccp4')
m = gemmi.read_ccp4_map(path)
self.assertEqual(m.grid.nu, 8)
self.assertEqual(m.grid.nv, 6)
self.assertEqual(m.grid.nw, 10)
self.assertEqual(m.header_i32(28), 0)
m.set_header_i32(28, 20140) # set NVERSION
self.assertEqual(m.header_i32(28), 20140)
dmax = m.header_float(21)
self.assertEqual(dmax, max(p.value for p in m.grid))
self.assertNotEqual(m.grid.axis_order, gemmi.AxisOrder.XYZ)
m.setup()
self.assertEqual(m.grid.axis_order, gemmi.AxisOrder.XYZ)
self.assertEqual(m.grid.nu, 60)
self.assertEqual(m.grid.nv, 24)
self.assertEqual(m.grid.nw, 60)
self.assertEqual(m.grid.point_count, 60 * 24 * 60)
self.assertEqual(m.header_float(14), 90.0) # 14 - alpha angle
self.assertEqual(m.grid.unit_cell.alpha, 90.0)
self.assertEqual(m.grid.spacegroup.ccp4, 4) # P21
pos = gemmi.Position(19.4, 3., 21.)
frac = m.grid.unit_cell.fractionalize(pos)
pos_value = 2.1543798446655273
self.assertAlmostEqual(m.grid.interpolate_value(pos), pos_value)
self.assertAlmostEqual(m.grid.interpolate_value(frac), pos_value)
# this spacegroup has symop -x, y+1/2, -z
m.grid.set_value(60 - 3, 24 // 2 + 4, 60 - 5,
100) # image of (3, 4, 5)
self.assertEqual(m.grid.get_value(60 - 3, 24 // 2 + 4, 60 - 5), 100)
self.assertTrue(math.isnan(m.grid.get_value(3, 4, 5)))
m.grid.symmetrize_max()
self.assertEqual(m.grid.get_value(3, 4, 5), 100)
m.grid.set_value(3, 4, 5, float('nan'))
self.assertTrue(math.isnan(m.grid.get_value(3, 4, 5)))
m.grid.symmetrize_min()
self.assertEqual(m.grid.get_value(3, 4, 5), 100)
m.grid.set_value(60 - 3, 24 // 2 + 4, 60 - 5, float('nan'))
m.grid.symmetrize_max()
self.assertEqual(m.grid.get_value(60 - 3, 24 // 2 + 4, 60 - 5), 100)
if numpy:
arr = numpy.array(m.grid, copy=False)
self.assertEqual(arr.shape, (60, 24, 60))
self.assertEqual(arr[3][4][5], 100)
grid2 = gemmi.FloatGrid(arr)
self.assertTrue(
numpy.allclose(m.grid, grid2, atol=0.0, rtol=0,
equal_nan=True))
def test_new(self):
N = 24
m = gemmi.FloatGrid(N, N, N)
self.assertEqual(m.nu, N)
self.assertEqual(m.nv, N)
self.assertEqual(m.nw, N)
m.set_value(1, 2, 3, 1.0)
self.assertEqual(sum(m), 1.0)
m.space_group = gemmi.find_spacegroup_by_name('C2')
self.assertEqual(m.space_group.number, 5)
m.symmetrize_max()
self.assertEqual(sum(m), 4.0)
m.fill(2.0)
m.space_group = gemmi.find_spacegroup_by_name('P 62 2 2')
self.assertEqual(len(m.space_group.operations()), 12)
m.set_value(1, 2, 3, 0.0)
m.symmetrize_min()
self.assertEqual(sum(m), 2 * N * N * N - 2 * 12)
def get_masked_sample(_reference_grid, _grid, _reference_structure):
mask = gemmi.FloatGrid(_reference_grid.nu,
_reference_grid.nv,
_reference_grid.nw, )
mask.set_unit_cell(_reference_grid.unit_cell)
mask.spacegroup = gemmi.find_spacegroup_by_name("P 1")
for model in _reference_structure:
for chain in model:
for residue in chain.get_polymer():
for atom in residue:
mask.set_points_around(atom.pos, 3.0, 1.0)
_reference_sample = np.array(_reference_grid, copy=False)
_sample = np.array(_grid, copy=False)
_mask = np.array(mask, copy=False)
sample_flattened = _sample[_mask != 0]
return sample_flattened
def get_cut_out_event_map(event_map: gemmi.FloatGrid,
coord: Coord,
radius: float = 10.0) -> gemmi.FloatGrid:
event_centroid = gemmi.Position(coord.x, coord.y, coord.z)
xmap_array = np.array(event_map, copy=True)
mask_grid = gemmi.Int8Grid(*xmap_array.shape)
# print(f"Spacegroup: {mask_grid.spacegroup.xhm()}")
# mask_grid.spacegroup = gemmi.find_spacegroup_by_name("P 21 21 21") # gemmi.find_spacegroup_by_name("P 1")#event_map.spacegroup
mask_grid.spacegroup = gemmi.find_spacegroup_by_name(
"P 1") #event_map.spacegroup
print(f"Spacegroup: {mask_grid.spacegroup.xhm()}")
print(f"grid: {mask_grid}")
mask_grid_array = np.array(mask_grid)
print(f"Mask grid array: {mask_grid_array.shape}")
print(f"Mask grid array: {mask_grid_array.size}")
print(f"Mask grid array: {np.sum(np.isfinite(mask_grid_array))}")
# print(f"Grid size: {mask_grid.size}")
mask_grid.set_unit_cell(event_map.unit_cell)
mask_grid.set_points_around(
event_centroid,
radius=radius,
value=1,
)
mask_grid.symmetrize_max()
mask_array = np.array(mask_grid, copy=False, dtype=np.int8)
new_grid = gemmi.FloatGrid(*xmap_array.shape)
new_grid.spacegroup = event_map.spacegroup # gemmi.find_spacegroup_by_name("P 1")
new_grid.set_unit_cell(event_map.unit_cell)
new_grid_array = np.array(new_grid, copy=False)
new_grid_array[np.nonzero(mask_array)] = xmap_array[np.nonzero(mask_array)]
new_grid.symmetrize_max()
return new_grid
def __setstate__(self, state):
data = state["data"]
self.grid = gemmi.FloatGrid(
data.shape[0],
data.shape[1],
data.shape[2],
)
spacegroup = state["spacegroup"]
self.grid.spacegroup = gemmi.SpaceGroup(spacegroup)
unit_cell = state["unit_cell"]
self.grid.unit_cell = gemmi.UnitCell(
unit_cell[0],
unit_cell[1],
unit_cell[2],
unit_cell[3],
unit_cell[4],
unit_cell[5],
)
for index, val in np.ndenumerate(data):
self.grid.set_value(index[0], index[1], index[2], data[index])
def score_structure(structure, xmap):
unit_cell = xmap.unit_cell
mask = gemmi.FloatGrid(xmap.nu, xmap.nv, xmap.nw)
mask.set_unit_cell(unit_cell)
mask.spacegroup = gemmi.find_spacegroup_by_name("P 1")
for model in structure:
for chain in model:
for residue in chain:
for atom_1 in residue:
if atom_1.element.name == "H":
continue
pos_1 = atom_1.pos
for atom_2 in residue:
if atom_2.element.name == "H":
continue
pos_2 = atom_2.pos
if pos_1.dist(pos_2) < 2.0:
new_pos = gemmi.Position(
(pos_1.x + pos_2.x) / 2,
(pos_1.y + pos_2.y) / 2,
(pos_1.z + pos_2.z) / 2,
)
mask.set_points_around(new_pos, 0.75, 1.0)
mask_array = np.array(mask)
xmap_array = np.array(xmap)
truncated_xmap_mask = xmap_array > 1.25
score = np.sum(truncated_xmap_mask * mask_array)
return float(score)
def get_mask(dataset: Dataset, alignment, sample_region, structure_factors,
sample_rate, ):
# Get the unaligned xmap
reflections: gemmi.Mtz = dataset.reflections
unaligned_xmap: gemmi.FloatGrid = reflections.transform_f_phi_to_map(
structure_factors.f,
structure_factors.phi,
sample_rate=sample_rate,
)
unaligned_xmap_array = np.array(unaligned_xmap, copy=False)
std = np.std(unaligned_xmap_array)
unaligned_xmap_array[:, :, :] = unaligned_xmap_array[:, :, :] / std
# Transform sample region
transformed_sample_region = get_transformed_sample_region(
sample_region,
alignment,
)
# mask
mask = gemmi.FloatGrid(unaligned_xmap.nu,
unaligned_xmap.nv,
unaligned_xmap.nw, )
mask.set_unit_cell(unaligned_xmap.unit_cell)
mask.spacegroup = gemmi.find_spacegroup_by_name("P 1")
for model in dataset.structure:
for chain in model:
for residue in chain.get_polymer():
for atom in residue:
mask.set_points_around(atom.pos, 3.0, 1.0)
# Get sample
# reference_sample = sample_dataset_global(mask, transformed_sample_region)
return mask
def get_event_map(event_map_file: Path) -> gemmi.FloatGrid:
m = gemmi.read_ccp4_map(str(event_map_file))
print(m.grid.spacegroup.xhm())
# m.grid.spacegroup = gemmi.find_spacegroup_by_name("P 21 21 21")
m.grid.spacegroup = gemmi.find_spacegroup_by_name("P 1")
print(m.grid.spacegroup.xhm())
m.setup()
grid_array = np.array(m.grid, copy=True)
print(m.grid.spacegroup.xhm())
print(m.grid)
print(grid_array.shape)
new_grid = gemmi.FloatGrid(*grid_array.shape)
new_grid.spacegroup = m.grid.spacegroup # gemmi.find_spacegroup_by_name("P 1")
new_grid.set_unit_cell(m.grid.unit_cell)
new_grid_array = np.array(new_grid, copy=False)
new_grid_array[:, :, :] = grid_array[:, :, :]
return new_grid
def clone_grid(grid):
new_grid = gemmi.FloatGrid(grid.nu, grid.nv, grid.nw)
new_grid.spacegroup = grid.spacegroup
new_grid.set_unit_cell(grid.unit_cell)
return new_grid
def resample(
reference_xmap: gemmi.FloatGrid,
moving_xmap: gemmi.FloatGrid,
reference_structure: gemmi.Structure,
moving_structure: gemmi.Structure,
monomerized=False,
):
# Get transform: from ref to align
transform = get_alignment(moving_structure, reference_structure, monomerized=monomerized)
print(f"Transform: {transform}; {transform.transform.vec} {transform.transform.mat}")
interpolated_grid = gemmi.FloatGrid(
reference_xmap.nu,
reference_xmap.nv,
reference_xmap.nw,
)
interpolated_grid.set_unit_cell(reference_xmap.unit_cell)
interpolated_grid.spacegroup = reference_xmap.spacegroup
# points
mask = gemmi.FloatGrid(reference_xmap.nu,
reference_xmap.nv,
reference_xmap.nw, )
mask.set_unit_cell(reference_xmap.unit_cell)
mask.spacegroup = gemmi.find_spacegroup_by_name("P 1")
for model in reference_structure:
for chain in model:
for residue in chain.get_polymer():
for atom in residue:
mask.set_points_around(atom.pos, 3.0, 1.0)
mask_array = np.array(mask)
mask_indicies = np.hstack([x.reshape((len(x), 1)) for x in np.nonzero(mask)])
print(f"Mask indicies shape: {mask_indicies.shape}")
fractional_coords = []
for model in reference_structure:
for chain in model:
for residue in chain.get_polymer():
for atom in residue:
fractional = reference_xmap.unit_cell.fractionalize(atom.pos)
fractional_coords.append([fractional.x, fractional.y, fractional.z])
fractional_coords_array = np.array(fractional_coords)
max_coord = np.max(fractional_coords_array, axis=0)
min_coord = np.min(fractional_coords_array, axis=0)
min_index = np.floor(min_coord * np.array([interpolated_grid.nu, interpolated_grid.nv, interpolated_grid.nw]))
max_index = np.floor(max_coord * np.array([interpolated_grid.nu, interpolated_grid.nv, interpolated_grid.nw]))
points = itertools.product(range(int(min_index[0]), int(max_index[0])),
range(int(min_index[1]), int(max_index[1])),
range(int(min_index[2]), int(max_index[2])),
)
# Unpack the points, poitions and transforms
point_list: List[Tuple[int, int, int]] = []
position_list: List[Tuple[float, float, float]] = []
transform_list: List[gemmi.transform] = []
com_moving_list: List[np.array] = []
com_reference_list: List[np.array] = []
transform_rotate_reference_to_moving = transform.transform
transform_rotate_reference_to_moving.vec.fromlist([0.0, 0.0, 0.0])
transform_reference_to_centered = gemmi.Transform()
transform_reference_to_centered.vec.fromlist((-transform.com_reference).tolist())
transform_reference_to_centered.mat.fromlist(np.eye(3).tolist())
tranform_centered_to_moving = gemmi.Transform()
tranform_centered_to_moving.vec.fromlist(transform.com_moving.tolist())
tranform_centered_to_moving.mat.fromlist(np.eye(3).tolist())
# indicies to positions
for point in points:
if mask.get_value(*point) < 1.0:
continue
# get position
position = interpolated_grid.get_position(*point)
# Tranform to origin frame
position_origin_reference = gemmi.Position(transform_reference_to_centered.apply(position))
# Rotate
position_origin_moving = gemmi.Position(transform_rotate_reference_to_moving.apply(position_origin_reference))
# Transform to moving frame
position_moving = gemmi.Position(tranform_centered_to_moving.apply(position_origin_moving))
# Interpolate moving map
interpolated_map_value = moving_xmap.interpolate_value(position_moving)
# Set original point
interpolated_grid.set_value(point[0], point[1], point[2], interpolated_map_value)
interpolated_grid.symmetrize_max()
return interpolated_grid
