def get_path(d_graph, src, max_weight, deltas):
# graph is a dict indexed by location
# nodes in graph must contain .weight members
edge = [pathnode(src, 0, coord.Coord(x=0, y=0))]
path = {src: pathnode(src, 0, coord.Coord(x=0, y=0))}
while len(edge) > 0:
edge.sort(key=lambda val: val.weight)
temp = edge.pop(0)
for delta in deltas:
loc = temp.loc + delta
if not loc in d_graph:
continue
elif loc in path:
continue
else:
for pnode in edge:
if pnode.loc == loc:
i = edge.index(loc)
if d_graph[loc].weight + temp.weight < edge[i].weight:
edge[i].weight = d_graph[loc].weight + temp.weight
break
else:
if d_graph[loc].weight + temp.weight <= max_weight:
edge.append(
pathnode(loc, d_graph[loc].weight + temp.weight,
delta))
if temp.weight <= max_weight:
path[temp.loc] = temp
return path
def consider_corners(data, head, board):
best = util.bad_move()
best_distance = 1
for corner in [[1, 1], [1, data['height'] - 2], [data['width'] - 2, 1],
[data['width'] - 2, data['height'] - 2]]:
path = board.path(coord.Coord(head[0], head[1]),
coord.Coord(corner[0], corner[1]))
distance = util.dist(head, corner)
if (path.cost - distance * DISTANCE_IS_GOOD_MULTIPLIER) < (
best.cost - best_distance *
DISTANCE_IS_GOOD_MULTIPLIER) and distance >= best_distance:
best = path
best_distance = distance
return best
def food(data, head, heatmap, board):
best = None
cost = 9996
for snack in data['food']:
path = board.path(coord.Coord(head[0], head[1]),
coord.Coord(snack[0], snack[1]))
heat = path.cost
if heat < cost:
best = path
cost = heat
if best is None or not util.is_valid_move(best, data):
return util.bad_move()
return best
def __set_opening(self): max_rank = self._team.get_ranks() position = self.__ground_coord for i in reversed(range(max_rank)): for j in range(self._team.get_num_rankers(i)): self.__opening_positions[(i, j)] = position position = coord.Coord(position.get_x() + self.__x_offset, position.get_y() + self.__y_offset)
def get_intersection(ray, segment):
r_point, r_dx, r_dy = ray.get_parametric_form()
r_px = r_point.get_x()
r_py = r_point.get_y()
r_mag = ray.get_magnitude()
s_point, s_dx, s_dy = segment.get_parametric_form()
s_px = s_point.get_x()
s_py = s_point.get_y()
s_mag = segment.get_magnitude()
if ray.get_slope() == segment.get_slope():
return None
t2 = ((r_dx * (s_py - r_py) + r_dy *
(r_px - s_px)) * 1.0) / (s_dx * r_dy - s_dy * r_dx)
if r_dx != 0:
t1 = ((s_px + s_dx * t2 - r_px) * 1.0) / r_dx
else:
t1 = ((s_py + s_dy * t2 - r_py) * 1.0) / r_dy
if t1 < 0:
return None
if t2 < 0 or t2 > 1:
return None
i_x = r_px + r_dx * t1
i_y = r_py + r_dy * t1
intersect_point = coord.Coord(i_x, i_y)
# print('Intersect point:')
# print(intersect_point)
return intersect_point, t1
def test_loc_init(self):
"Test the Coord object creation with loc"
# 1. Create Area object from text
myobj = coord.Coord(x=1, y=5)
# 2. Make sure it has the expected values
self.assertEqual(myobj.x, 1)
self.assertEqual(myobj.y, 5)
self.assertEqual(myobj.isInfinite, None)
self.assertEqual(myobj.minX, None)
self.assertEqual(myobj.maxX, None)
self.assertEqual(myobj.minY, None)
self.assertEqual(myobj.maxY, None)
self.assertEqual(myobj.area, None)
self.assertEqual(myobj.closest, None)
# 3. Test the distance method
self.assertEqual(myobj.distance(1, 5), 0)
self.assertEqual(myobj.distance(3, 5), 2)
self.assertEqual(myobj.distance(0, 5), 1)
self.assertEqual(myobj.distance(1, 7), 2)
self.assertEqual(myobj.distance(1, 4), 1)
self.assertEqual(myobj.distance(0, 0), 6)
self.assertEqual(myobj.distance(9, 9), 12)
def __init__(self, agent_type, team, map_manager):
super(LineOpeningSkill, self).__init__(agent_type, team, map_manager)
self.__opening_positions = {}
self.__openings_created = False
self.__x_offset = 5
self.__y_offset = 0
self.__ground_coord = coord.Coord(5, 5)
def wiggle(data, board, head, dist=5):
best = util.bad_move()
for x in range(-dist, dist + 1):
target_x = head[0] + x
if target_x < 0 or target_x >= data['width']: continue
for y in range(-dist, dist + 1):
target_y = head[1] + y
if target_y < 0 or target_y >= data['height']: continue
if (abs(x) + abs(y)) == dist:
path = board.path(coord.Coord(head[0], head[1]),
coord.Coord(target_x, target_y))
if path.cost < best.cost and util.is_valid_move(path, data):
best = path
return best
def __init__(self):
random.seed(None)
self.window_size = WINDOW_SIZE
self.map_size = MAP_SIZE
self.size = coord.Coord(x=1.0 / self.map_size.x,
y=1.0 / self.map_size.y)
#self.tiles = {}
#for y in range(self.map_size.y):
# for x in range(self.map_size.x):
# ttex = 0
# tx = x if TILE_TYPE == tile.RECT else (x + (0.5 if y%2 == 0 else 0.0))
# ty = y
# tloc = coord.Coord(x=tx, y=ty)
# tweight = random.randint(RANDRANGE[0], RANDRANGE[1])
# #self.tiles.append(tile.tile(tx, ty, ttex, TILE_TYPE, tweight))
# self.tiles[tloc] = tile.tile(tx, ty, ttex, TILE_TYPE, tweight)
coordFactory = map.RectCoordFactory(MAP_SIZE)
dataFactory = map.RandDataFactory(RANDRANGE)
self.map = map.Map(coordFactory, dataFactory)
self.init_window()
self.init_callback()
tex.init()
self.units = {
coord.Coord(x=0, y=0):
unit.unit(coord.Coord(x=0, y=0), g_texnames.index("unit"),
2), # * sum(RANDRANGE)/2),
coord.Coord(x=4, y=6):
unit.unit(coord.Coord(x=4, y=6), g_texnames.index("unit"),
4), # * sum(RANDRANGE)/2),
coord.Coord(x=3, y=6):
unit.unit(coord.Coord(x=3, y=6), g_texnames.index("unit"),
1), # * sum(RANDRANGE)/2),
coord.Coord(x=3, y=5):
unit.unit(coord.Coord(x=3, y=5), g_texnames.index("unit"),
0), # * sum(RANDRANGE)/2)
}
self.selected = None
self.mouseloc = None
self.wmouseloc = None
self.tooltip = tooltips.tooltip(1.0)
self.tooltip.start()
self.path = None
self.pathtex = None
self.myfps = fps.fps()
def mouse(self, button, state, x, y):
# TODO: ADD HEX CLICK LOGIC.
if TILE_TYPE != tile.RECT: return
rx = x * self.map_size.x // self.window_size.x
ry = (self.window_size.y - y) * self.map_size.y // self.window_size.y
ri = ry * self.map_size.x + rx
if button == GLUT_RIGHT_BUTTON and state == GLUT_DOWN:
if coord.Coord(x=rx, y=ry) in self.units:
self.selected = coord.Coord(x=rx, y=ry)
else:
self.selected = None
elif button == GLUT_LEFT_BUTTON and state == GLUT_DOWN:
c = coord.Coord(x=rx, y=ry)
if self.path != None and c in self.path and not c in self.units:
self.units[c] = self.units[self.selected]
self.units.pop(self.selected, None)
self.selected = c
self.units[self.selected].loc = self.selected
self.update_path()
def mouse_passive(self, x, y):
self.wmouseloc = coord.Coord(x=x, y=y)
self.wmouseloc *= coord.Coord(x=1.0, y=-1.0)
self.wmouseloc += coord.Coord(x=0, y=self.window_size.y - 1)
self.wmouseloc /= self.window_size
nx = x * self.map_size.x // self.window_size.x
ny = (self.window_size.y - y) * self.map_size.y // self.window_size.y
self.mouseloc = coord.Coord(x=nx, y=ny)
self.update_path()
#if self.mouseloc in self.tiles:
#self.tooltip.data = self.tiles[self.mouseloc].weight
#self.tooltip.start()
if self.mouseloc in self.map:
self.tooltip.data = self.map[self.mouseloc].weight
self.tooltip.start()
def follow(data, head, board):
target = False
for snake in data['snakes']:
snake_tail = snake['coords'][-1]
if util.dist(head, snake_tail) == 1:
target = snake_tail
if not target:
# No tails nearby
return util.bad_move()
path = board.path(coord.Coord(head[0], head[1]),
coord.Coord(target[0], target[1]))
if path is None or not util.is_valid_move(path, data):
return util.bad_move()
return path
def __set_opening(self): gamemap = self._map_manager.get_map() max_rank = self._team.get_ranks() for i in reversed(range(max_rank)): for j in range(self._team.get_num_rankers(i)): found_position = False while not found_position: position = coord.Coord(random.randint(0, gamemap.get_map_width()), random.randint(0, gamemap.get_map_height())) if not self.__check_within_obstacle(position) and not self.__check_already_occupied(position): self.__opening_positions[(i, j)] = position found_position = True
def get_parametric_form(self):
'''
Returns point and direction
'''
p_x = self.get_a().get_x()
p_y = self.get_a().get_y()
point = coord.Coord(p_x, p_y)
d_x = self.get_b().get_x() - self.get_a().get_x()
d_y = self.get_b().get_y() - self.get_a().get_y()
# assert(d_x != d_y)
return point, d_x, d_y
def get_mid_edge_points(self, epsilon=0):
'''
A mid-edge point is defined as a point which lies on the middle of one of the edges of a rectangle. Since a
a rectangle has 4 edges, it will have 4 such mid-edge points
epsilon: Distance by which a mid-edge point should be distant from the edge
'''
mid_edge_points = []
w2 = self.__width / 2 + epsilon
h2 = self.__height / 2 + epsilon
factor = [-1, 1]
for fa in factor:
a = self.__centre[0] + w2 * fa
b = self.__centre[1]
mid_edge_points.append(coord.Coord(a, b))
a = self.__centre[0]
b = self.__centre[1] + h2 * fa
mid_edge_points.append(coord.Coord(a, b))
return mid_edge_points
def consider_our_tail(board, data, head):
oursnake = data['oursnake']['coords']
if (len(oursnake) == 0):
return util.bad_move() # didn't find our snake, bail
target = oursnake[-1]
if (target == head):
if head[1] > 1:
target = [head[0], head[1] - 1]
else:
target = [head[0], head[1] + 1]
if target in oursnake[:-1]:
return util.bad_move()
path = board.path(coord.Coord(head[0], head[1]),
coord.Coord(target[0], target[1]))
if path is None or not util.is_valid_move(path, data):
return util.bad_move()
return path
def get_visibility_polygon(self, current_position, current_rotation, num_rays, visibility_angle):
# c = coord.Coord(self.x, self.y)
vis_points = [int(current_position.get_x()), int(current_position.get_y())]
rotation = current_rotation - visibility_angle
offset = (visibility_angle * 2.0)/num_rays
bbox = self.get_bbox(current_position)
hits = list(self.__rtree_idx.intersection(bbox.get_rtree_bbox(), objects=True))
polygon_hits = [item.object for item in hits]
# polygon_ids = [item.id for item in hits]
# print('Polygons considered:', polygon_ids)
nearby_polygons = polygon_hits + [bbox] + [self.__boundary_polygon]
while rotation < current_rotation + visibility_angle:
rotation_x = math.cos(coord.Coord.to_radians(-rotation))
rotation_y = math.sin(coord.Coord.to_radians(-rotation))
r = coord.Coord(current_position.get_x() + rotation_x, current_position.get_y() + rotation_y)
rotation += offset
if r.get_x() < 0 or r.get_x() > self.__width or r.get_y() < 0 or r.get_y() > self.__height:
vis_points.append(int(current_position.get_x()))
vis_points.append(int(current_position.get_y()))
continue
ray = shapes.Line(current_position, r)
# print('ray:', ray)
closest_intersect = None
for polygon in nearby_polygons:
# print('polygon:',polygon)
for i in range(polygon.get_num_lines()):
intersect = shapes.Line.get_intersection(ray, polygon.get_line(i))
if not intersect:
continue
if not closest_intersect or intersect[1] < closest_intersect[1]:
closest_intersect = intersect
if not closest_intersect:
print('Closest intersect not found')
print('From coordinate:', current_position)
print('Ray:', ray)
print('Segment:', polygon.get_line(i))
continue
vis_points.append(int(closest_intersect[0].get_x()))
vis_points.append(int(closest_intersect[0].get_y()))
vis_points_tuple = tuple(vis_points)
visibility_polygon = shapes.Polygon(vis_points_tuple)
return visibility_polygon
def __init__(self,
points_tuple,
offset=1.0,
line_analysis=False,
point_analysis=False):
assert (len(points_tuple) % 2 == 0)
# print('')
# print('Points tuple:', points_tuple)
self.__num_vertices = len(points_tuple) // 2
self.__vertices = []
self.__line_analysis = line_analysis
self.__lines = []
self.__point_analysis = point_analysis
self.__mpl_path = None
self.__points_tuple = points_tuple
if offset == 1.0:
i = 0
while i < self.__num_vertices * 2:
vertex = coord.Coord(points_tuple[i], points_tuple[i + 1])
self.__vertices.append(vertex)
i += 2
else:
i = 0
offsetted_list = []
while i < self.__num_vertices * 2:
vertex = (points_tuple[i], points_tuple[i + 1])
offsetted_list.append(vertex)
i += 2
offsetted_tuple = tuple(offsetted_list)
pco = pyclipper.PyclipperOffset()
pco.AddPath(offsetted_tuple, pyclipper.JT_ROUND,
pyclipper.ET_CLOSEDPOLYGON)
solution = pco.Execute(offset)[0]
for vertex in solution:
vcoord = coord.Coord(vertex[0], vertex[1])
self.__vertices.append(vcoord)
def test_empty_init(self):
"Test the default Coord creation"
# 1. Create default Area object
myobj = coord.Coord()
# 2. Make sure it has the default values
self.assertEqual(myobj.x, None)
self.assertEqual(myobj.y, None)
self.assertEqual(myobj.isInfinite, None)
self.assertEqual(myobj.minX, None)
self.assertEqual(myobj.maxX, None)
self.assertEqual(myobj.minY, None)
self.assertEqual(myobj.maxY, None)
self.assertEqual(myobj.area, None)
self.assertEqual(myobj.closest, None)
def processText(self, text):
# 1. Loop for all of the lines in the text
for line in text:
# 2. Get the x and y coordinates
parts = line.split(',')
x=int(parts[0])
y=int(parts[1])
# 3. Create and save the coordinate
self.coords.append(coord.Coord(x=x, y=y))
# 4. Keep track of the maximum x and y
if x > self.maxX:
self.maxX = x
if y > self.maxY:
self.maxY = y
def move_idle_dumb(data, head, heatmap, board):
left_pathdata = board.path(coord.Coord(head[0], head[1]),
coord.Coord(max(0, head[0] - 1), head[1]))
right_pathdata = board.path(
coord.Coord(head[0], head[1]),
coord.Coord(min(len(heatmap) - 1, head[0] + 1), head[1]))
up_pathdata = board.path(coord.Coord(head[0], head[1]),
coord.Coord(head[0], max(0, head[1] - 1)))
down_pathdata = board.path(
coord.Coord(head[0], head[1]),
coord.Coord(head[0], min(len(heatmap[0]) - 1, head[1] + 1)))
smallest = min(left_pathdata.cost, right_pathdata.cost, up_pathdata.cost,
down_pathdata.cost)
if smallest == left_pathdata.cost:
return left_pathdata
if smallest == right_pathdata.cost:
return right_pathdata
if smallest == up_pathdata.cost:
return up_pathdata
if smallest == down_pathdata.cost:
return down_pathdata
def totalDimensions(self):
# 1. Create a coordinate to hold the dimensions
coordinate = coord.Coord(x=self.maxX // 2, y=self.maxY // 2)
# 2. Check that it satisifies
if self.totalDistance(coordinate.x, coordinate.y) >= self.total:
print("(%d, %d) does not satisify" % (coordinate.x, coordinate.y))
return None
# 2. Find the lowest X dimension
for coordinate.minX in range(coordinate.x, -1, -1):
if self.totalDistance(coordinate.minX, coordinate.y) > self.total:
coordinate.minX += 1
break
# 3. Find the lowest y dimension
for coordinate.minY in range(coordinate.y, -1, -1):
if self.totalDistance(coordinate.x, coordinate.minY) > self.total:
coordinate.minY += 1
break
# 4. Find the highest X dimension
for coordinate.maxX in range(coordinate.x, self.maxX+1):
if self.totalDistance(coordinate.maxX, coordinate.y) > self.total:
coordinate.maxX -= 1
break
# 5. Find the highest y dimension
for coordinate.maxY in range(coordinate.y, self.maxY+1):
if self.totalDistance(coordinate.x, coordinate.maxY) > self.total:
coordinate.maxY -= 1
break
# 6. Determine if infinite or finite
if coordinate.minX == 0 or coordinate.minY == 0 or coordinate.maxX == self.maxX or coordinate.maxY == self.maxY:
coordinate.isInfinite = True
else:
coordinate.isInfinite = False
# 7. return coordinate with dimensions
return coordinate
def run():
pygame.init()
pygame.display.gl_set_attribute(pygame.GL_CONTEXT_MAJOR_VERSION, 4)
pygame.display.gl_set_attribute(pygame.GL_CONTEXT_MINOR_VERSION, 1)
pygame.display.gl_set_attribute(pygame.GL_CONTEXT_PROFILE_MASK,
pygame.GL_CONTEXT_PROFILE_CORE)
pygame.display.set_mode((1280, 720), pygame.DOUBLEBUF | pygame.OPENGL)
program = get_program()
groups = obj_loader.load('shangwu', 'part1.obj')
env_obj = env.Env()
arrow_obj = arrow.Arrow()
coord_obj = coord.Coord()
uniform.get_locs(program)
camera_obj = camera.Camera()
running = True
clock = pygame.time.Clock()
while running:
clock.tick(50)
for event in pygame.event.get():
if event.type == pygame.QUIT:
running = False
camera_obj.process_event(event)
uniform.set_view(camera_obj.view)
draw(program, groups)
env_obj.draw()
coord_obj.draw()
pygame.display.flip()
def __init__(self, graph, nodeList):
# Todo: "unreachable"
if PATHFINDING_DEBUG: print "Nodelist:", nodeList
self._nodes = nodeList
self.cost = 0
if len(nodeList) < 2:
self.cost = 9992
self.nextCoord = None
self.nextDirection = "No"
print "Path of length 1, looks impossible?"
return
self.length = len(nodeList)
self.nextCoord = coord.Coord(nodeList[1][0], nodeList[1][1])
self.nextDirection = "????"
if self.nextCoord.x == nodeList[0][0]:
if self.nextCoord.y < nodeList[0][1]:
self.nextDirection = "up"
else:
self.nextDirection = "down"
elif self.nextCoord.y == nodeList[0][1]:
if self.nextCoord.x > nodeList[0][0]:
self.nextDirection = "right"
else:
self.nextDirection = "left"
weights = nx.get_edge_attributes(graph, 'weight')
for nodeSeq in range(0, len(nodeList) - 1):
cur = nodeList[nodeSeq]
next = nodeList[nodeSeq + 1]
if (cur, next) in weights:
self.cost += weights[(cur, next)]
else:
self.cost += weights[(next, cur)]
self.cost -= constants.ourHeadWeight
if PATHFINDING_DEBUG: print "cost:", self.cost
def get_target_regions(args, chrom, words):
"""Parse start and end positions and return list of Coord "
objects representing arget region(s)."""
start_words = words[7].split(";")
end_words = words[8].split(";")
if len(start_words) != len(end_words):
raise coord.CoordError("number of start (%d) and end (%d) positions "
"do not match" %
(len(start_words), len(end_words)))
n_coord = len(start_words)
region_list = []
for i in range(n_coord):
start = int(start_words[i])
end = int(end_words[i])
region = coord.Coord(chrom, start, end)
if args.target_region_size:
if region.length() != args.target_region_size:
# override the size of the target region
# with size provided on command line
mid = (region.start + region.end) / 2
region.start = mid - args.target_region_size / 2
region.end = mid + args.target_region_size / 2
if region.start < 1:
region.start = 1
if region.end > chrom.length:
region.end = chrom.length
region_list.append(region)
return region_list
import unit
import tooltips
import random
import coord
import pathfinding
import tex
import fps
import geometry
import map
# TODO: next step, implement turn structure.
window = 0 # glut window number
TILE_TYPE = tile.RECT
MAP_SIZE = coord.Coord(x=10, y=10)
WINDOW_SIZE = coord.Coord(x=640, y=480)
FULLSCREEN = False
USEFONT = False
TEXDIR = "textures/"
TEXEXT = ".png"
RANDRANGE = (1, 3)
SHOWFPS = True
TILE_ADJ = [
coord.Coord(x=-1, y=0),
coord.Coord(x=0, y=-1),
coord.Coord(x=1, y=0),
coord.Coord(x=0, y=1)
]
def get_current_coordinate(self):
return coord.Coord(self.__x, self.__y)
import coord
WINDOW_SIZE = coord.Coord(x=640, y=480)
window = 0 # glut window number
def refresh2d(vw, vh, width, height):
glViewport(0, 0, width, height)
glClearColor(0, 0, 0, 0)
glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
glLoadIdentity()
glMatrixMode(GL_PROJECTION)
glLoadIdentity()
glOrtho(0.0, vw, 0.0, vh, 0.0, 1.0)
glMatrixMode(GL_MODELVIEW)
glLoadIdentity()
def run():
width = 1280
height = 720
pygame.init()
pygame.display.gl_set_attribute(pygame.GL_CONTEXT_MAJOR_VERSION, 4)
pygame.display.gl_set_attribute(pygame.GL_CONTEXT_MINOR_VERSION, 1)
pygame.display.gl_set_attribute(pygame.GL_CONTEXT_PROFILE_MASK,
pygame.GL_CONTEXT_PROFILE_CORE)
pygame.display.set_mode((width, height), pygame.DOUBLEBUF|pygame.OPENGL)
start_server(st)
program = get_program()
model = pyrr.matrix44.create_from_translation(pyrr.Vector3([0, 0, 0]))
projection = pyrr.matrix44.create_perspective_projection_matrix(
45, 1.0 * width / height, 0.1, 1000)
uniform.get_locs(program)
uniform.set_model(model)
uniform.set_projection(projection)
camera_obj = camera.Camera()
env_obj = env.Env()
coord_obj = coord.Coord()
model3d_obj = model3d.Model3d('shangwu', 'part1.obj')
running = True
clock = pygame.time.Clock()
while running:
for event in pygame.event.get():
if event.type == pygame.QUIT:
running = False
camera_obj.process_event(event)
glClearColor(0.0, 0.0, 0.0, 1.0)
glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
uniform.set_view(camera_obj.view)
env_obj.draw()
coord_obj.draw()
#model3d_obj.draw(model_q=g['fusion_obj'].mag_fusion_obj.q)
#model3d_obj.draw(model_q=g['fusion_obj'].adjusted_gyro_q)
model3d_obj.draw(model_q=g['fusion_obj'].q)
if g['fusion_obj'].adjusted_axis is not None:
axis = g['fusion_obj'].adjusted_axis
x_arrow = arrow.Arrow(color=[1.0, 0.0, 0.0],
vector=axis[0])
#x_arrow.draw()
y_arrow = arrow.Arrow(color=[0.0, 1.0, 0.0],
vector=axis[1])
#y_arrow.draw()
z_arrow = arrow.Arrow(color=[0.0, 0.0, 1.0],
vector=axis[2])
#z_arrow.draw()
if g['fusion_obj'].acc_normal is not None:
cycle_obj = cycle.Cycle(cos=g['fusion_obj'].acc_normal[0])
#cycle_obj.draw()
cycle_obj = cycle.Cycle(cos=g['fusion_obj'].mag_normal[0],
nv=g['fusion_obj'].mag_v)
#cycle_obj.draw()
pygame.display.flip()
def main():
args = parse_args()
write_header(sys.stdout)
# find index of individual in list of samples
ind_idx = lookup_individual_index(args, args.individual)
data_files = DataFiles(args)
chrom_list = chromosome.get_all_chromosomes(args.chrom)
chrom_dict = chromosome.get_chromosome_dict(args.chrom)
genomewide_read_counts = get_genomewide_count(data_files.read_count_h5,
chrom_list)
unknown_chrom = set([])
if util.is_gzipped(args.input_file):
f = gzip.open(args.input_file, "rt")
else:
f = open(args.input_file, "r")
line_count = 0
if args.target_region_size:
sys.stderr.write("setting target region size to %d\n" %
args.target_region_size)
for line in f:
line_count += 1
if line_count % 1000 == 0:
sys.stderr.write(".")
if line.startswith("#"):
continue
words = line.rstrip().split()
if words[1] == "NA":
# no SNP defined on this line:
write_NA_line(sys.stdout)
continue
chrom_name = words[0]
if chrom_name in chrom_dict:
chrom = chrom_dict[chrom_name]
else:
if not chrom_name.startswith("chr"):
# try adding 'chr' to front of name
new_chrom_name = "chr" + chrom_name
if new_chrom_name in chrom_dict:
chrom_name = new_chrom_name
chrom = chrom_dict[chrom_name]
else:
# can't figure out this chromosome name
if not chrom_name in unknown_chrom:
unknown_chrom.add(chrom_name)
sys.stderr.write("WARNING: unknown chromosome '%s'")
continue
region_list = get_target_regions(args, chrom, words)
snp_pos = int(words[1])
snp_ref_base = words[3]
snp_alt_base = words[4]
# TODO: check that SNP ref/alt match?
snp_region = coord.Coord(chrom, snp_pos, snp_pos)
# pull out all of the SNPs in the target region(s)
region_snps = get_region_snps(data_files, region_list, ind_idx)
# pull out test SNP
test_snp_list = get_region_snps(data_files, [snp_region], ind_idx)
if len(test_snp_list) != 1:
test_snp = None
sys.stderr.write("WARNING: could not find test SNP at "
"position %s:%d\n" % (chrom.name, snp_pos))
het_snps = []
else:
test_snp = test_snp_list[0]
# pull out haplotype counts from linked heterozygous SNPs
het_snps = get_het_snps(region_snps)
set_snp_counts(data_files, region_list, het_snps, test_snp, args)
region_read_counts = get_region_read_counts(data_files, region_list)
write_output(sys.stdout, region_list, het_snps, test_snp, snp_pos,
region_read_counts, genomewide_read_counts)
sys.stderr.write("\n")
f.close()
data_files.close()
